llava_qwen.py

#    Copyright 2024 Hao Zhang
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.
from typing import List, Optional, Tuple, Union, Dict
import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss
import transformers
from transformers import AutoConfig, AutoModelForCausalLM, LlamaConfig, LlamaModel, LlamaForCausalLM
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.generation.utils import GenerateOutput
# from .llava_arch import LlavaMetaModel, LlavaMetaForCausalLM
from .modeling_qwen2 import Qwen2Config, Qwen2Model, Qwen2ForCausalLM
# from transformers import Qwen2Config, Qwen2Model, Qwen2ForCausalLM
import pdb
import time
import random
random.seed(42)
import torch
from statistics import mean
import torch.nn.functional as F
import PIL
from decord import VideoReader, cpu
from .conversation import conv_templates, SeparatorStyle
from .constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, DEFAULT_IMAGE_TOKEN
from .mm_utils import tokenizer_image_token, load_video, KeywordsStoppingCriteria, get_anyres_image_grid_shape
import math
import re
from .vision_tower_builder import build_vision_tower
from .vision_resampler_builder import build_vision_resampler
from .vision_projector_builder import build_vision_projector
from .utils import rank0_print
from .sae import SiglipAE
import numpy as np
import pdb
from abc import ABC, abstractmethod

class LlavaMetaModel:
    def __init__(self, config):
        super(LlavaMetaModel, self).__init__(config)

        if hasattr(config, "mm_vision_tower"):
            delay_load = getattr(config, "delay_load", False)
            self.vision_tower = build_vision_tower(config, delay_load=delay_load)
            self.vision_resampler = build_vision_resampler(config, vision_tower=self.vision_tower)
            self.mm_projector = build_vision_projector(config, vision_cfg=self.vision_tower.config)

            if "unpad" in getattr(config, "mm_patch_merge_type", ""):
                self.image_newline = nn.Parameter(torch.empty(config.hidden_size, dtype=self.dtype))
    
        self.hidden_size=config.hidden_size
        self.text_mlp=nn.Sequential(
            nn.Linear(config.hidden_size,config.hidden_size),
            nn.GELU(),
        )
        self.sae=SiglipAE()

    def get_vision_tower(self):
        vision_tower = getattr(self, "vision_tower", None)
        if type(vision_tower) is list:
            vision_tower = vision_tower[0]
        return vision_tower

    def initialize_vision_modules(self, model_args, fsdp=None):
        vision_tower = model_args.vision_tower
        mm_vision_select_layer = model_args.mm_vision_select_layer
        mm_vision_select_feature = model_args.mm_vision_select_feature
        pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter
        mm_patch_merge_type = model_args.mm_patch_merge_type

        self.config.mm_vision_tower = vision_tower
        self.config.vision_tower_pretrained = getattr(model_args, "vision_tower_pretrained", "")

        if self.get_vision_tower() is None:
            vision_tower = build_vision_tower(model_args)
            vision_resampler = build_vision_resampler(model_args, vision_tower=vision_tower)
            for k, v in vision_resampler.config.items():
                setattr(self.config, k, v)

            if fsdp is not None and len(fsdp) > 0:
                self.vision_tower = [vision_tower]
                self.vision_resampler = [vision_resampler]
            else:
                self.vision_tower = vision_tower
                self.vision_resampler = vision_resampler
        else:
            if fsdp is not None and len(fsdp) > 0:
                vision_resampler = self.vision_resampler[0]
                vision_tower = self.vision_tower[0]
            else:
                vision_resampler = self.vision_resampler
                vision_tower = self.vision_tower
            vision_tower.load_model()

            # In case it is frozen by LoRA
            for p in self.vision_resampler.parameters():
                p.requires_grad = True

        self.config.use_mm_proj = True
        self.config.mm_projector_type = getattr(model_args, "mm_projector_type", "linear")
        self.config.mm_hidden_size = getattr(vision_resampler, "hidden_size", vision_tower.hidden_size)
        self.config.mm_vision_select_layer = mm_vision_select_layer
        self.config.mm_vision_select_feature = mm_vision_select_feature
        self.config.mm_patch_merge_type = mm_patch_merge_type
        
        self.sae=SiglipAE()
        self.sae.load_state_dict(torch.load('/share/LXRlxr0_0/code/videoxl2/videoxl2/longva/longva/model/encoder.pth'),strict=False)
        
        if getattr(self, "mm_projector", None) is None:
            self.mm_projector = build_vision_projector(self.config, vision_cfg=vision_tower.config)

            if "unpad" in mm_patch_merge_type:
                embed_std = 1 / torch.sqrt(torch.tensor(self.config.hidden_size, dtype=self.dtype))
                self.image_newline = nn.Parameter(torch.randn(self.config.hidden_size, dtype=self.dtype) * embed_std)
        else:
            # In case it is frozen by LoRA
            for p in self.mm_projector.parameters():
                p.requires_grad = True

        if pretrain_mm_mlp_adapter is not None:
            mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location="cpu")

            def get_w(weights, keyword):
                return {k.split(keyword + ".")[1]: v for k, v in weights.items() if keyword in k}

            incompatible_keys = self.mm_projector.load_state_dict(get_w(mm_projector_weights, "mm_projector"))
            rank0_print(f"Loaded mm projector weights from {pretrain_mm_mlp_adapter}. Incompatible keys: {incompatible_keys}")
            incompatible_keys = self.vision_resampler.load_state_dict(get_w(mm_projector_weights, "vision_resampler"), strict=False)
            rank0_print(f"Loaded vision resampler weights from {pretrain_mm_mlp_adapter}. Incompatible keys: {incompatible_keys}")
        
def unpad_image(tensor, original_size):
    """
    Unpads a PyTorch tensor of a padded and resized image.

    Args:
    tensor (torch.Tensor): The image tensor, assumed to be in CxHxW format.
    original_size (tuple): The original size of the image (height, width).

    Returns:
    torch.Tensor: The unpadded image tensor.
    """
    original_width, original_height = original_size
    current_height, current_width = tensor.shape[1:]

    # Compute aspect ratios
    original_aspect_ratio = original_width / original_height
    current_aspect_ratio = current_width / current_height

    # Determine padding size and direction
    if original_aspect_ratio > current_aspect_ratio:
        # Padding was added to the height
        scale_factor = current_width / original_width
        new_height = int(original_height * scale_factor)
        padding = (current_height - new_height) // 2
        unpadded_tensor = tensor[:, padding : current_height - padding, :]
    else:
        # Padding was added to the width
        scale_factor = current_height / original_height
        new_width = int(original_width * scale_factor)
        padding = (current_width - new_width) // 2
        unpadded_tensor = tensor[:, :, padding : current_width - padding]

    return unpadded_tensor

class LlavaMetaForCausalLM(ABC):
    @abstractmethod
    def get_model(self):
        pass

    def get_vision_tower(self):
        return self.get_model().get_vision_tower()

    def get_2dPool(self, image_feature):
        height = width = self.get_vision_tower().num_patches_per_side
        num_frames, num_tokens, num_dim = image_feature.shape
        image_feature = image_feature.view(num_frames, height, width, -1)
        image_feature = image_feature.permute(0, 3, 1, 2).contiguous()
        # image_feature = nn.functional.max_pool2d(image_feature, self.config.mm_spatial_pool_stride)
        if self.config.mm_spatial_pool_mode == "average":
            image_feature = nn.functional.avg_pool2d(image_feature, self.config.mm_spatial_pool_stride)
        elif self.config.mm_spatial_pool_mode == "max":
            image_feature = nn.functional.max_pool2d(image_feature, self.config.mm_spatial_pool_stride)
        else:
            raise ValueError(f"Unexpected mm_spatial_pool_mode: {self.config.mm_spatial_pool_mode}")
        image_feature = image_feature.permute(0, 2, 3, 1)
        image_feature = image_feature.view(num_frames, -1, num_dim)
        return image_feature

    def encode_images(self, images):
        image_features = self.get_model().get_vision_tower()(images)
        #image_features = self.get_model().vision_resampler(image_features, images=images)
        image_features = self.get_model().mm_projector(image_features)
        image_features = self.get_model().vision_resampler(image_features, images=images)
        return image_features

    def add_image(self, image_features):
        return torch.repeat_interleave(image_features, repeats=4, dim=0)
    
    def add_video(self, video_features):
        # Current batch size
        current_batch_size = video_features.size(0)
        # Handle cases where the batch size is less than 4
        if current_batch_size < 4:
            last_feature = video_features[-1:]
            # Calculate how many times the last feature needs to be repeated
            num_repeats = 4 - current_batch_size
            repeated_features = last_feature.repeat(num_repeats, 1, 1, 1)
            # Concatenate original features with repeated last feature
            expanded_x = torch.cat([video_features, repeated_features], dim=0)
            return expanded_x

        # Handle cases where the batch size is 4 or greater, but not a multiple of 4
        if current_batch_size % 4 != 0:
            last_feature = video_features[-1:]
            # Calculate how many features are needed to reach the next multiple of 4
            padding_size = 4 - (current_batch_size % 4)
            repeated_features = last_feature.repeat(padding_size, 1, 1, 1)
            # Concatenate original features with repeated last feature
            expanded_x = torch.cat([video_features, repeated_features], dim=0)
            return expanded_x
        
        # If the batch size is already a multiple of 4, return as is
        return video_features
    def encode_multimodals(self, videos_or_images, video_idx_in_batch, split_sizes=None):
        if self.config.enable_chunk_prefill:
            chunk_size_for_vision_tower = self.config.prefill_config['chunk_size_for_vision_tower']
        else:
            chunk_size_for_vision_tower = 100000
        # pdb.set_trace()
        # Define the maximum batch size (1024 frames)
        max_batch_size = chunk_size_for_vision_tower
        # print(f'max_batch_size: {max_batch_size}')
        num_frames = videos_or_images.shape[0]
        # Initialize a list to store the features from each batch
        videos_or_images_features = []

        videos_or_images_features = torch.empty((num_frames, 729, 1152), device=self.get_model().device, dtype=self.get_model().dtype)
       
        # Split videos_or_images into smaller batches if num_frames > max_batch_size
        current_idx = 0
        if num_frames > max_batch_size:
            # Calculate the number of batches needed
            num_batches = (num_frames + max_batch_size - 1) // max_batch_size
            for i in range(num_batches):
                start_idx = i * max_batch_size
                end_idx = min((i + 1) * max_batch_size, num_frames)

                # Process each batch separately
                batch_videos_or_images = videos_or_images[start_idx:end_idx]
                batch_features = self.get_model().get_vision_tower()(batch_videos_or_images)
                # videos_or_images_features.append(batch_features)

                videos_or_images_features[current_idx:current_idx + batch_features.shape[0]] = batch_features
                # Update the current index for the next batch
                current_idx += batch_features.shape[0]
                # peak_memory_allocated = torch.cuda.max_memory_allocated()
                # print(f"vision encoder 显存峰值: {peak_memory_allocated / (1024**3):.2f} GB") # 转换为GB
            
            # Concatenate the features of all batches
            # videos_or_images_features = torch.cat(videos_or_images_features, dim=0)
        else:
            videos_or_images_features = self.get_model().get_vision_tower()(videos_or_images)

        per_videos_or_images_features = torch.split(videos_or_images_features, split_sizes, dim=0) 
        all_videos_or_images_features = []

        # peak_memory_allocated = torch.cuda.max_memory_allocated()
        # print(f"vision encoder 显存峰值: {peak_memory_allocated / (1024**3):.2f} GB") # 转换为GB
        del videos_or_images_features
        torch.cuda.empty_cache()

        chunk_size = chunk_size_for_vision_tower
        # print(f'chunk_size: {chunk_size}')
        all_feat_list = []
        for idx, feat in enumerate(per_videos_or_images_features):
            for i in range(0, feat.shape[0], chunk_size):
                batched_feat = feat[i:i+chunk_size] # chunk_size = 48, batched_feat.shape=[48, 729, 1152]
                batched_feat=self.interpolate(batched_feat) # 插值后 batched_feat.shape=[48, 1152, 24, 24]
                if idx in video_idx_in_batch:
                    batched_feat = self.add_video(batched_feat) # 第一纬度补充到4的倍数
                else:
                    batched_feat = self.add_image(batched_feat)

                bc,ch,h,w = batched_feat.shape
                batched_feat = batched_feat.view(bc//4,ch,4,h,w)

                batched_feat = self.get_model().sae(batched_feat).squeeze(2)
                batched_feat = batched_feat.permute(0, 2, 3, 1).contiguous().flatten(1, 2)
                
                batched_feat = self.get_model().mm_projector(batched_feat)
                batched_feat = self.get_2dPool(batched_feat)
                all_feat_list.append(batched_feat)

        feat = torch.cat(all_feat_list, dim=0)
        # peak_memory_allocated = torch.cuda.max_memory_allocated()
        # print(f"sae 显存峰值: {peak_memory_allocated / (1024**3):.2f} GB") # 转换为GB

        del per_videos_or_images_features
        del all_feat_list
        torch.cuda.empty_cache()
        all_videos_or_images_features.append(feat)
        return all_videos_or_images_features
   
    def interpolate(self,image_features):
        b, num_tokens, dim = image_features.shape
        
        #print(str(image_features.shape)+' i\n')
        
        target_h = target_w = int(576**0.5)
        h = w = int(num_tokens**0.5)

        image_features = image_features.view(b, h, w, dim)
        image_features = image_features.permute(0, 3, 1, 2).contiguous()

        chunk_size = 24
        chunks = torch.split(image_features, chunk_size, dim=0)
        interpolated_chunks = []
        for chunk in chunks:
            interpolated_chunk = F.interpolate(
                chunk.to(torch.float32),
                size=(target_h, target_w),
                mode="bilinear",
                align_corners=False,
            ).to(chunk.dtype)
            interpolated_chunks.append(interpolated_chunk)
        image_features = torch.cat(interpolated_chunks, dim=0)
        del interpolated_chunks
        
        del chunks

        return image_features

    def prepare_inputs_labels_for_multimodal(self, input_ids, position_ids, attention_mask, past_key_values, labels, images, modalities=["image"], image_sizes=None,time_embedding=None):

        vision_tower = self.get_vision_tower()
        if vision_tower is None or images is None or input_ids.shape[1] == 1:
            return input_ids, position_ids, attention_mask, past_key_values, None, labels

        if type(images) is list or images.ndim == 5:
            if type(images) is list:
                images = [x.unsqueeze(0) if x.ndim == 3 else x for x in images]

            video_idx_in_batch = []
            for _ in range(len(modalities)):
                if modalities[_] == "video":
                    video_idx_in_batch.append(_)

            images_list = []
            for image in images:
                if image.ndim == 4:
                    images_list.append(image)
                else:
                    images_list.append(image.unsqueeze(0))
            #print(len(images_list),images_list[0].shape)

            concat_images = torch.cat([image for image in images_list], dim=0)
            split_sizes = [image.shape[0] for image in images_list]

            image_features = self.encode_multimodals(concat_images, video_idx_in_batch, split_sizes)    #16,144,3584

            mm_patch_merge_type = getattr(self.config, "mm_patch_merge_type", "flat")
            image_aspect_ratio = getattr(self.config, "image_aspect_ratio", "square")

            visual_drop_score=[]
            new_image_features=[]
            
            if mm_patch_merge_type == "flat":
                
                if image_features[0].ndim>2:
                    image_features = [x.flatten(0, 1) for x in image_features]
            elif mm_patch_merge_type== "unires":
                #print('unires')
                for image_idx, image_feature in enumerate(image_features):
                    # rank0_print(f"Initial feature size : {image_feature.shape}")
                    if image_idx in video_idx_in_batch:  # video operations
                        #print(image_feature.shape)
                        image_feature = image_feature.flatten(0, 1)
                        
                    elif image_feature.shape[0] > 1:
                        # base image feature is never used in unires
                        base_image_feature = image_feature[0]
                        image_feature = image_feature[1:]

                        height = width = self.get_vision_tower().num_patches_per_side
                        assert height * width == base_image_feature.shape[0]

                        kernel_size = mm_patch_merge_type.split("avgpool")[-1].split("x")[-1]
                        kernel_size = 2
                        image_feature = image_feature.view(image_feature.shape[0], height, width, -1) # [4, 24, 24, 4096]
                        image_feature = image_feature.permute(0, 3, 1, 2).contiguous() # [4, 4096, 24, 24]
                        image_feature = nn.functional.avg_pool2d(image_feature,kernel_size) # [4, 4096, 12, 12]
                        image_feature = image_feature.flatten(2, 3) # [4, 4096, 144]
                        image_feature = image_feature.permute(0, 2, 1).contiguous() # [4, 144, 4096]
                        
                        #print(image_feature.shape)
                        image_feature = image_feature.flatten(0, 1)
                        
                    else:
                        
                        image_feature = image_feature[0]
                        
                    new_image_features.append(image_feature)
                
                image_features = new_image_features
                
            elif mm_patch_merge_type.startswith("spatial"):
                new_image_features = []
                for image_idx, image_feature in enumerate(image_features):
                    # FIXME: now assume the image is square, and split to 2x2 patches
                    # num_patches = h * w, where h = w = sqrt(num_patches)
                    # currently image_feature is a tensor of shape (4, num_patches, hidden_size)
                    # we want to first unflatten it to (2, 2, h, w, hidden_size)
                    if image_idx in video_idx_in_batch:  # video operations
                        if "unpad" in mm_patch_merge_type:
                            # image_feature = image_feature.permute(2, 0, 1).contiguous()
                            # image_feature =  torch.cat((image_feature, self.model.image_newline[:, None, None].expand(*image_feature.shape[:-1], 1).to(image_feature.device)), dim=-1)
                            # image_feature = image_feature.permute(1, 2, 0).contiguous()
                            image_feature = image_feature.flatten(0, 1)
                            image_feature = torch.cat((image_feature, self.model.image_newline[None].to(image_feature.device)), dim=0)

                    elif image_feature.shape[0] > 1:  # multi patches and multi images operations
                        base_image_feature = image_feature[0]
                        image_feature = image_feature[1:]
                        height = width = self.get_vision_tower().num_patches_per_side
                        assert height * width == base_image_feature.shape[0]

                        if "anyres_max" in image_aspect_ratio:
                            matched_anyres_max_num_patches = re.match(r"anyres_max_(\d+)", image_aspect_ratio)
                            if matched_anyres_max_num_patches:
                                max_num_patches = int(matched_anyres_max_num_patches.group(1))

                        if image_aspect_ratio == "anyres" or "anyres_max" in image_aspect_ratio:
                            if hasattr(self.get_vision_tower(), "image_size"):
                                vision_tower_image_size = self.get_vision_tower().image_size
                            else:
                                raise ValueError("vision_tower_image_size is not found in the vision tower.")
                            num_patch_width, num_patch_height = get_anyres_image_grid_shape(image_sizes[image_idx], self.config.image_grid_pinpoints, vision_tower_image_size)
                            image_feature = image_feature.view(num_patch_height, num_patch_width, height, width, -1)
                        else:
                            image_feature = image_feature.view(2, 2, height, width, -1)

                        if "maxpool2x2" in mm_patch_merge_type:
                            image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous()
                            image_feature = image_feature.flatten(1, 2).flatten(2, 3)
                            image_feature = nn.functional.max_pool2d(image_feature, 2)
                            image_feature = image_feature.flatten(1, 2).transpose(0, 1)
                        elif "unpad" in mm_patch_merge_type and "anyres_max" in image_aspect_ratio and matched_anyres_max_num_patches:
                            unit = image_feature.shape[2]
                            image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous()
                            image_feature = image_feature.flatten(1, 2).flatten(2, 3)
                            image_feature = unpad_image(image_feature, image_sizes[image_idx])
                            c, h, w = image_feature.shape
                            times = math.sqrt(h * w / (max_num_patches * unit**2))
                            if times > 1.1:
                                image_feature = image_feature[None]
                                image_feature = nn.functional.interpolate(image_feature, [int(h // times), int(w // times)], mode="bilinear")[0]
                            image_feature = torch.cat((image_feature, self.model.image_newline[:, None, None].expand(*image_feature.shape[:-1], 1).to(image_feature.device)), dim=-1)
                            image_feature = image_feature.flatten(1, 2).transpose(0, 1)
                        elif "unpad" in mm_patch_merge_type:
                            image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous()
                            image_feature = image_feature.flatten(1, 2).flatten(2, 3)
                            image_feature = unpad_image(image_feature, image_sizes[image_idx])
                            image_feature = torch.cat((image_feature, self.model.image_newline[:, None, None].expand(*image_feature.shape[:-1], 1).to(image_feature.device)), dim=-1)
                            image_feature = image_feature.flatten(1, 2).transpose(0, 1)
                        else:
                            image_feature = image_feature.permute(0, 2, 1, 3, 4).contiguous()
                            image_feature = image_feature.flatten(0, 3)
                        if "nobase" in mm_patch_merge_type:
                            pass
                        else:
                            image_feature = torch.cat((base_image_feature, image_feature), dim=0)
                    else:  # single image operations
                        image_feature = image_feature[0]
                        if "unpad" in mm_patch_merge_type:
                            image_feature = torch.cat((image_feature, self.model.image_newline[None]), dim=0)

                    new_image_features.append(image_feature)
                image_features = new_image_features
            else:
                raise ValueError(f"Unexpected mm_patch_merge_type: {self.config.mm_patch_merge_type}")
        else:
            error_message = """
            Something is wrong with the input shape. Most likely, you did not wrap the image or video input in a list:
            This is correct:
                model.generate(input_ids, images=[video_tensor],  modalities=["video"], **gen_kwargs)
                model.generate(input_ids, images=[image_tensor],  modalities=["image"], **gen_kwargs)
            This is wrong:
                model.generate(input_ids, images=video_tensor,  modalities=["video"], **gen_kwargs)
                model.generate(input_ids, images=image_tensor,  modalities=["image"], **gen_kwargs)
            """
            raise ValueError(error_message)

        #print(time_embedding[0].shape)
        #video_token_indices=[]
        for image_idx, image_feature in enumerate(image_features):
            if time_embedding[image_idx] is not None:
                mask = (time_embedding[image_idx] == 151654)
                indices = torch.nonzero(mask).squeeze()

                embed_token=self.get_model().embed_tokens(time_embedding[image_idx])
                embed_token[indices]=image_features[image_idx]
                
                #video_token_indices.append(indices)
                
                image_features[image_idx]=embed_token
    
        if getattr(self.config, "tune_mm_mlp_adapter", False) and getattr(self.config, "mm_use_im_start_end", False):
            raise NotImplementedError

        # Let's just add dummy tensors if they do not exist,
        # it is a headache to deal with None all the time.
        # But it is not ideal, and if you have a better idea,
        # please open an issue / submit a PR, thanks.
        _labels = labels
        _position_ids = position_ids
        _attention_mask = attention_mask
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
        else:
            attention_mask = attention_mask.bool()
        if position_ids is None:
            position_ids = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=input_ids.device)
        if labels is None:
            labels = torch.full_like(input_ids, IGNORE_INDEX)

        # remove the padding using attention_mask -- FIXME
        _input_ids = input_ids
        input_ids = [cur_input_ids[cur_attention_mask] for cur_input_ids, cur_attention_mask in zip(input_ids, attention_mask)]
        labels = [cur_labels[cur_attention_mask] for cur_labels, cur_attention_mask in zip(labels, attention_mask)]

        new_input_embeds = []
        new_labels = []
        cur_image_idx = 0

        for batch_idx, cur_input_ids in enumerate(input_ids):
            num_images = (cur_input_ids == IMAGE_TOKEN_INDEX).sum()
            #print(num_images)
            if num_images>=2:
                print(num_images,input_ids)
            if num_images == 0:
                cur_image_features = image_features[cur_image_idx]
                cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids)
                cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0]], dim=0)
                new_input_embeds.append(cur_input_embeds)
                new_labels.append(labels[batch_idx])
                cur_image_idx += 1
                continue

            image_token_indices = [-1] + torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0].tolist() + [cur_input_ids.shape[0]]
            #print(image_token_indices) #[-1, 14, 236]
            cur_input_ids_noim = []
            cur_labels = labels[batch_idx]
            
            # print(cur_input_ids)
            # print(labels[batch_idx])
            
            cur_labels_noim = []
            for i in range(len(image_token_indices) - 1):
                cur_input_ids_noim.append(cur_input_ids[image_token_indices[i] + 1 : image_token_indices[i + 1]])
                cur_labels_noim.append(cur_labels[image_token_indices[i] + 1 : image_token_indices[i + 1]])
            split_sizes = [x.shape[0] for x in cur_labels_noim]
            
           #print(torch.cat(cur_input_ids_noim).shape,torch.cat(cur_input_ids_noim))
            
            cur_input_embeds = self.get_model().embed_tokens(torch.cat(cur_input_ids_noim))
            cur_input_embeds_no_im = torch.split(cur_input_embeds, split_sizes, dim=0)
            cur_new_input_embeds = []
            cur_new_labels = []

            for i in range(num_images + 1):
                cur_new_input_embeds.append(cur_input_embeds_no_im[i])
                cur_new_labels.append(cur_labels_noim[i])
                if i < num_images:
                    ##############
                    cur_image_features = image_features[cur_image_idx]
                    cur_image_idx += 1
                    cur_new_input_embeds.append(cur_image_features)
                    cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=cur_labels.device, dtype=cur_labels.dtype))

            cur_new_input_embeds = [x.to(self.device) for x in cur_new_input_embeds]

            # import pdb; pdb.set_trace()
            cur_new_input_embeds = torch.cat(cur_new_input_embeds)

            cur_new_labels = torch.cat(cur_new_labels)

            new_input_embeds.append(cur_new_input_embeds)
            new_labels.append(cur_new_labels)

        # Truncate sequences to max length as image embeddings can make the sequence longer
        tokenizer_model_max_length = getattr(self.config, "tokenizer_model_max_length", None)
        # NOTE: qmh 
        # new_input_embeds = [x[:tokenizer_model_max_length] for x, modality in zip(new_input_embeds, modalities)]
        # new_labels = [x[:tokenizer_model_max_length] for x, modality in zip(new_labels, modalities)]

        # TODO: Hard code for control loss spike
        # if tokenizer_model_max_length is not None:
        #     new_input_embeds = [x[:4096] if modality != "video" else x[:tokenizer_model_max_length] for x, modality in zip(new_input_embeds, modalities)]
        #     new_labels = [x[:4096] if modality != "video" else x[:tokenizer_model_max_length] for x, modality in zip(new_labels, modalities)]

        # Combine them
        max_len = max(x.shape[0] for x in new_input_embeds)
        batch_size = len(new_input_embeds)

        new_input_embeds_padded = []
        new_labels_padded = torch.full((batch_size, max_len), IGNORE_INDEX, dtype=new_labels[0].dtype, device=new_labels[0].device)
        attention_mask = torch.zeros((batch_size, max_len), dtype=attention_mask.dtype, device=attention_mask.device)
        position_ids = torch.zeros((batch_size, max_len), dtype=position_ids.dtype, device=position_ids.device)

        for i, (cur_new_embed, cur_new_labels) in enumerate(zip(new_input_embeds, new_labels)):
            cur_len = cur_new_embed.shape[0]
            if getattr(self.config, "tokenizer_padding_side", "right") == "left":
                new_input_embeds_padded.append(torch.cat((torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device), cur_new_embed), dim=0))
                if cur_len > 0:
                    new_labels_padded[i, -cur_len:] = cur_new_labels
                    attention_mask[i, -cur_len:] = True
                    position_ids[i, -cur_len:] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)
            else:
                new_input_embeds_padded.append(torch.cat((cur_new_embed, torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)), dim=0))
                if cur_len > 0:
                    new_labels_padded[i, :cur_len] = cur_new_labels
                    attention_mask[i, :cur_len] = True
                    position_ids[i, :cur_len] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)

        new_input_embeds = torch.stack(new_input_embeds_padded, dim=0)

        if _labels is None:
            new_labels = None
        else:
            new_labels = new_labels_padded

        if _attention_mask is None:
            attention_mask = None
        else:
            attention_mask = attention_mask.to(dtype=_attention_mask.dtype)

        if _position_ids is None:
            position_ids = None
        if getattr(self.config, "use_pos_skipping", False) and self.training:
            position_ids = torch.arange(new_input_embeds.size(1), device=new_input_embeds.device).unsqueeze(0).to(new_input_embeds.device)
            split_position = random.randint(0, new_input_embeds.size(1))
            left_add = random.randint(0, self.config.pos_skipping_range)
            right_add = random.randint(left_add, self.config.pos_skipping_range)
            position_ids[:, :split_position] += left_add
            position_ids[:, split_position:] += right_add
        # import pdb; pdb.set_trace()
        return None, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels

    def initialize_vision_tokenizer(self, model_args, tokenizer):
        if model_args.mm_use_im_patch_token:
            tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
            self.resize_token_embeddings(len(tokenizer))

        if model_args.mm_use_im_start_end:
            num_new_tokens = tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
            self.resize_token_embeddings(len(tokenizer))

            if num_new_tokens > 0:
                input_embeddings = self.get_input_embeddings().weight.data
                output_embeddings = self.get_output_embeddings().weight.data

                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)

                input_embeddings[-num_new_tokens:] = input_embeddings_avg
                output_embeddings[-num_new_tokens:] = output_embeddings_avg

            if model_args.tune_mm_mlp_adapter:
                for p in self.get_input_embeddings().parameters():
                    p.requires_grad = True
                for p in self.get_output_embeddings().parameters():
                    p.requires_grad = False

            if model_args.pretrain_mm_mlp_adapter:
                mm_projector_weights = torch.load(model_args.pretrain_mm_mlp_adapter, map_location="cpu")
                embed_tokens_weight = mm_projector_weights["model.embed_tokens.weight"]
                assert num_new_tokens == 2
                if input_embeddings.shape == embed_tokens_weight.shape:
                    input_embeddings[-num_new_tokens:] = embed_tokens_weight[-num_new_tokens:]
                elif embed_tokens_weight.shape[0] == num_new_tokens:
                    input_embeddings[-num_new_tokens:] = embed_tokens_weight
                else:
                    raise ValueError(f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}.")

        elif model_args.mm_use_im_patch_token:
            if model_args.tune_mm_mlp_adapter:
                for p in self.get_input_embeddings().parameters():
                    p.requires_grad = False
                for p in self.get_output_embeddings().parameters():
                    p.requires_grad = False


class LlavaQwenConfig(Qwen2Config):
    model_type = "llava_qwen"


class LlavaQwenModel(LlavaMetaModel, Qwen2Model):
    config_class = LlavaQwenConfig

    def __init__(self, config: Qwen2Config):
        super(LlavaQwenModel, self).__init__(config)


class LlavaQwenForCausalLM(Qwen2ForCausalLM, LlavaMetaForCausalLM):
    config_class = LlavaQwenConfig

    def __init__(self, config):
        # super(Qwen2ForCausalLM, self).__init__(config)
        Qwen2ForCausalLM.__init__(self, config)
        config.model_type = "llava_qwen"
        config.rope_scaling = None

        self.model = LlavaQwenModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        # Initialize weights and apply final processing
        self.post_init()

    def get_model(self):
        return self.model

    def uniform_sampling(self, embeds, start_idx, end_idx, step):
        indices = torch.arange(start_idx, end_idx, step).to(device=embeds.device)
        return embeds.index_select(1, indices), indices
    def pooling_sampling(self, embeds, start_idx, end_idx, step, pool_type='avg'):
        selected = embeds[:, start_idx:end_idx, :]
        B, D, L = selected.shape
        kernel_size = step
        stride = step
        
        selected_transposed = selected.transpose(1, 2)  # shape: (1, 12, 4)

        if pool_type == 'avg_pool':
            pooled = F.avg_pool1d(selected_transposed, kernel_size=kernel_size, stride=stride)
        elif pool_type == 'max_pool':
            pooled = F.max_pool1d(selected_transposed, kernel_size=kernel_size, stride=stride)
        else:
            raise ValueError(f"Unsupported pooling type: {pool_type}")

        pooled = pooled.transpose(1, 2)  # shape: (1, 2, 12)
        return pooled, torch.arange(start_idx, start_idx + pooled.shape[1] * step, step).to(device=embeds.device)

    def process_block(self, block_embeds, current_past_key_values=None, bsz=1, device=None, position_ids=None, key_position_ids=None):
        if current_past_key_values is None:
            seq_len = block_embeds.size(1)
            position_ids = torch.arange(0, seq_len, device=device).expand(bsz, -1)
            attention_mask = torch.ones((bsz, seq_len), device=device, dtype=torch.long)
        else:
            seq_len = block_embeds.size(1)
            prefix_len = current_past_key_values[0][0].size(2)
            attention_mask = torch.ones((bsz, prefix_len + seq_len), device=device, dtype=torch.long)


        outputs = self.model(
            inputs_embeds=block_embeds,
            attention_mask=attention_mask,
            position_ids=position_ids,
            key_position_ids=key_position_ids,
            past_key_values=current_past_key_values,
            use_cache=True,
            return_dict=True,
        )
        return outputs.past_key_values

    def pooling_kvs(self, kvs, step):
        # kvs shape: (bsz, 4, seq_len, head_dim)
        kernel_size = step
        stride = step
        # kvs = kvs.transpose(2, 3)
        # pooled_kvs = F.avg_pool1d(kvs, kernel_size=kernel_size, stride=stride)
        kvs_permuted = kvs.permute(0, 1, 3, 2) # (batch_size, num_heads, feature_dim, sequence_length)
        N_flat = kvs_permuted.shape[0] * kvs_permuted.shape[1]
        C = kvs_permuted.shape[2]
        L = kvs_permuted.shape[3]
        kvs_for_pool = kvs_permuted.reshape(N_flat, C, L)
        pooled_kvs = F.avg_pool1d(kvs_for_pool, kernel_size=kernel_size, stride=stride)
        pooled_kvs_restored = pooled_kvs.view(kvs.shape[0], kvs.shape[1], pooled_kvs.shape[1], pooled_kvs.shape[2]).permute(0, 1, 3, 2)
        return pooled_kvs_restored


    def get_sparse_attention_mask(self, total_len, num_blocks, block_size, time_token_start_indices, time_token_end_indices, time_token_indices, visual_token_start_pos, visual_token_end_pos, attention_mask, inputs_embeds, prev_blocks_num=None):

        causal_mask = torch.tril(torch.ones((total_len, total_len), dtype=torch.bool)).unsqueeze(0).repeat(1, 1, 1)
        mask = torch.zeros(total_len, total_len, dtype=torch.bool)
        start = visual_token_start_pos

        record_block_start = []
        for i in range(num_blocks):
            next_time_token_pos = (i + 1)*block_size
            if next_time_token_pos >= len(time_token_start_indices):
                end = visual_token_end_pos
            else:
                end =  time_token_start_indices[ next_time_token_pos ]

            mask[start:end, start:end] = True  

            if len(record_block_start) >= prev_blocks_num:
                prev_start = record_block_start[-prev_blocks_num]
            else:
                prev_start = visual_token_start_pos

            mask[start:end, prev_start:start] = True   
            record_block_start.append(start)
            start = end

            
        mask[:, :visual_token_start_pos] = True
        mask[visual_token_end_pos:, :] = True        
     
        for idx in time_token_indices:
            mask[idx, :] = True   
            mask[:, idx] = True   

        causal_mask = torch.tril(torch.ones(total_len, total_len, dtype=torch.bool))
        final_mask = (mask & causal_mask).unsqueeze(0).unsqueeze(0).to(dtype=attention_mask.dtype, device=attention_mask.device)

        num_allowed = final_mask.sum().item()
        upper_triangle_num = total_len * (total_len + 1) // 2
        ratio = num_allowed / upper_triangle_num
        
        invert_mask = 1.0 - final_mask
        final_mask = ((1.0 - final_mask) * -1e9).to(dtype=inputs_embeds.dtype)
        return final_mask, ratio


    def cat_history_kvs(self, prefix_kvs, kvs_part2, kvs_part3):
        prefix_kvs = [[kvs] for kvs in prefix_kvs]
        cat_kvs = []
        for prefix_kvs_this_layer, kvs_part2_this_layer, kvs_part3_this_layer in zip(prefix_kvs, kvs_part2, kvs_part3):
            prefix_key_this_layer = [tmp[0] for tmp in prefix_kvs_this_layer]
            prefix_val_this_layer = [tmp[1] for tmp in prefix_kvs_this_layer]

            key_part2_this_layer = [tmp[0] for tmp in kvs_part2_this_layer]
            val_part2_this_layer = [tmp[1] for tmp in kvs_part2_this_layer]

            key_part3_this_layer = [tmp[0] for tmp in kvs_part3_this_layer]
            val_part3_this_layer = [tmp[1] for tmp in kvs_part3_this_layer]

            key_this_layer = torch.cat(prefix_key_this_layer + key_part2_this_layer + key_part3_this_layer, dim=-2)
            val_this_layer = torch.cat(prefix_val_this_layer + val_part2_this_layer + val_part3_this_layer, dim=-2)

            cat_kvs.append((key_this_layer, val_this_layer))
        return cat_kvs

    def forward_streaming(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        key_position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        dpo_forward: Optional[bool] = False,
        cache_position=None,
        visual_token_start_pos=None,
        visual_token_end_pos=None,
        time_token_start_indices=None,
        frames_num=None,
        time_token_indices=None,
        time_token_end_indices=None,
        block_size_chosed=None,
        prev_blocks_num=None,
        offload: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        
        block_size = block_size_chosed 
        visual_token_start_pos = visual_token_start_pos
        visual_token_end_pos = visual_token_end_pos
        visual_len = visual_token_end_pos - visual_token_start_pos
        num_blocks = (frames_num  + block_size * 4 - 1) // (block_size * 4)
        
        # streaming inps
        blocks_positions = [[(0, 0, visual_token_start_pos)]]
        frames_groups = [(0, visual_token_start_pos)]
        for idx, (time_start, time_end) in enumerate(zip(time_token_start_indices, time_token_end_indices)):
            if idx + 1 < len(time_token_start_indices):
                frames_group_end = time_token_start_indices[idx + 1]
            else:
                frames_group_end = visual_token_end_pos
            frames_groups.append(
                (time_start, time_end, frames_group_end)
            )

        single_block = []
        for group in frames_groups[1:]: 
            single_block.append(group)
            if len(single_block) == block_size:
                blocks_positions.append(single_block)
                single_block = []
        if len(single_block) != 0:
           blocks_positions.append(single_block)
        num_blocks = len(blocks_positions)

        start = time.time()
        record_prefill_time = 0

        full_inputs_embeds = inputs_embeds 
        bsz, total_len, embed_dim = full_inputs_embeds.size()
        device = full_inputs_embeds.device

        prefix_embeds = full_inputs_embeds[:, :visual_token_start_pos, :]
        visual_embeds = full_inputs_embeds[:, visual_token_start_pos:visual_token_end_pos, :]
        suffix_embeds = full_inputs_embeds[:, visual_token_end_pos:, :]
        num_visual_tokens = visual_embeds.size(1)

        all_past_key_values = [[] for _ in range(len(self.model.layers))] 
        prefix_past_key_values = []  

        # torch.cuda.reset_peak_memory_stats()
        
        if prefix_embeds.size(1) > 0:
            pkv = self.process_block(prefix_embeds, bsz=bsz, device=device)
            for i in range(len(pkv)):
                all_past_key_values[i].append(pkv[i])
                prefix_past_key_values.append(pkv[i])

        prev_blocks = blocks_positions[1:1+prev_blocks_num]
        prev_the_first_block = prev_blocks[0]
        prev_b_start = prev_the_first_block[0][0]
        prev_the_last_block = prev_blocks[-1]
        prev_b_end = prev_the_last_block[-1][-1]

        block_streaming_past_key_values = prefix_past_key_values

        query_position_ids = torch.arange(prev_b_start, prev_b_end, dtype=torch.long, device=device)
        past_key_position_ids = torch.arange(0, block_streaming_past_key_values[0][0].size(2), dtype=torch.long, device=device)
        key_position_ids = torch.cat([past_key_position_ids, query_position_ids], dim=0)

        visual_embeds_this_block = full_inputs_embeds[:,prev_b_start:prev_b_end,:]
        pkv = self.process_block(visual_embeds_this_block, current_past_key_values=block_streaming_past_key_values, bsz=bsz, device=device, position_ids=query_position_ids.unsqueeze(0), key_position_ids=key_position_ids.unsqueeze(0))

        for i in range(len(pkv)):
            for block in prev_blocks:
                block_start, _, _ = block[0]
                _, _, block_end = block[-1]
                all_past_key_values[i].append( (pkv[i][0][:,:,block_start:block_end], pkv[i][1][:,:,block_start:block_end]) )

        block_streaming_past_key_values_part1 = prefix_past_key_values
        position_ids_part1 = torch.arange(0, prefix_past_key_values[0][0].size(2), dtype=torch.long, device=device)
        block_streaming_past_key_values_part2 = [[] for _ in range(len(self.model.layers))]  
        position_ids_part2 = torch.tensor([], dtype=torch.long, device=device)
        block_streaming_past_key_values_part3=None
        position_ids_part3 = None
        
        query_position_ids = None
        for idx, single_block in enumerate(blocks_positions[:]):
            
            if idx == 0 or idx <= prev_blocks_num:
                continue
            
            b_start, _, _ = single_block[0]
            _, _, b_end = single_block[-1]
            visual_embeds_this_block = full_inputs_embeds[:,b_start:b_end,:]
            prev_blocks = blocks_positions[max(idx - prev_blocks_num, 1):idx]
            prev_the_first_block = prev_blocks[0]
            prev_b_start = prev_the_first_block[0][0]
            
            this_block_length = b_end - prev_b_start
            prev_block_length = b_start - prev_b_start
            true_block_length = b_end - b_start

            block_streaming_past_key_values_part3 = [tmp[-prev_blocks_num:] for tmp in all_past_key_values]

            if offload:
                block_streaming_past_key_values_part3 = [
                    [
                        (t[0].to(device=device), t[1].to(device=device))
                        for t in sublist
                    ]
                    for sublist in block_streaming_past_key_values_part3
                ]

            block_streaming_past_key_values = self.cat_history_kvs(block_streaming_past_key_values_part1, block_streaming_past_key_values_part2, block_streaming_past_key_values_part3)

            query_position_ids = torch.arange(b_start, b_end, dtype=torch.long, device=device)
            position_ids_part3 = torch.arange(prev_b_start, b_start, dtype=torch.long, device=device)
            key_position_ids = torch.cat([position_ids_part1, position_ids_part2, position_ids_part3, query_position_ids], dim=0)

            start_1 = time.time() 
            pkv = self.process_block(visual_embeds_this_block, current_past_key_values=block_streaming_past_key_values, bsz=bsz, device=device, position_ids=query_position_ids.unsqueeze(0), key_position_ids=key_position_ids.unsqueeze(0))
            end_1 = time.time()

            record_prefill_time += end_1-start_1

            for i in range(len(pkv)):
                length_before_chunk = block_streaming_past_key_values[i][0].size(2)
                key_this_block, val_this_block = pkv[i]
                key_this_block = key_this_block[:,:,length_before_chunk:,:]
                val_this_block = val_this_block[:,:,length_before_chunk:,:]
                
                if offload:
                    all_past_key_values[i].append( (key_this_block.to('cpu'), val_this_block.to('cpu')) )
                else:
                    all_past_key_values[i].append( (key_this_block, val_this_block) )

                time_keys_list = []
                time_vals_list = []

                extract_timestamps_position_ids_list = []
                for group in prev_the_first_block:
                    time_start, time_end, _ = group
                    extract_timestamps_position_ids_list.append(torch.arange(time_start, time_end, dtype=torch.long, device=device))

                    time_start = time_start - prev_b_start
                    time_end = time_end - prev_b_start

                    time_keys_list.append(block_streaming_past_key_values_part3[i][0][0][:,:,time_start:time_end,:])
                    time_vals_list.append(block_streaming_past_key_values_part3[i][0][1][:,:,time_start:time_end,:])
                    
                time_keys = torch.cat(time_keys_list, dim=2)
                time_vals = torch.cat(time_vals_list, dim=2)

                block_streaming_past_key_values_part2[i].append( (time_keys, time_vals) )

                if i == 0: 
                    position_ids_part2 = torch.cat([position_ids_part2] + extract_timestamps_position_ids_list, dim=0)


        merged_pkv = []
        for layer_pkvs in all_past_key_values:
            if not layer_pkvs:
                continue
            keys = torch.cat([pkv[0].to(device=device) for pkv in layer_pkvs], dim=2)  # dim=2 是 sequence 维度
            values = torch.cat([pkv[1].to(device=device) for pkv in layer_pkvs], dim=2)
            merged_pkv.append((keys, values))

        # peak_memory_allocated = torch.cuda.max_memory_allocated()
        # print(f"prefill 显存峰值: {peak_memory_allocated / (1024**3):.2f} GB") # 转换为GB
    
        pkv = merged_pkv
        del block_streaming_past_key_values
        del all_past_key_values
        del block_streaming_past_key_values_part1
        del block_streaming_past_key_values_part2
        del block_streaming_past_key_values_part3
        torch.cuda.empty_cache()

        # TODO: bi-decoding acceleration
        mixed_prefill_past_key_values = pkv
        prefill_len = visual_token_end_pos

        # torch.cuda.reset_peak_memory_stats()
        # Process suffix
        if suffix_embeds.size(1) > 0:
            seq_len = suffix_embeds.size(1)
            total_len = prefill_len + seq_len
            position_ids = torch.arange(prefill_len, total_len, device=device, dtype=torch.long).expand(bsz, -1)
            key_position_ids = torch.arange(0, total_len, device=device, dtype=torch.long).expand(bsz, -1)
            attention_mask = torch.ones((bsz, total_len), device=device, dtype=torch.long)

            outputs = super().forward(
                inputs_embeds=suffix_embeds,
                attention_mask=attention_mask,
                position_ids=position_ids,
                key_position_ids=key_position_ids,
                past_key_values=mixed_prefill_past_key_values,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                use_cache=True,
                return_dict=return_dict,
                # blocks_positions=None,
            )
            # peak_memory_allocated = torch.cuda.max_memory_allocated()
            # print(f"decoding 显存峰值: {peak_memory_allocated / (1024**3):.2f} GB") # 转换为GB
            del mixed_prefill_past_key_values
            torch.cuda.empty_cache()

        return outputs
    def forward_mask(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        dpo_forward: Optional[bool] = False,
        cache_position=None,
        visual_token_start_pos=None,
        visual_token_end_pos=None,
        time_token_start_indices=None,
        time_token_end_indices=None,
        frames_num=None,
        time_token_indices=None,
        prev_blocks_num=None,
        block_size_chosed=None
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        bsz, total_len, embed_dim = inputs_embeds.size()
        visual_token_start_pos = visual_token_start_pos
        visual_token_end_pos = visual_token_end_pos
        visual_len = visual_token_end_pos - visual_token_start_pos
        
        block_size_list = [2,4,8,16,32]
        best_block_size = None
        min_diff = float('inf')
       
        block_size = block_size_chosed
        num_blocks = (frames_num  + block_size * 4 - 1) // (block_size * 4)
        final_mask, ratio = self.get_sparse_attention_mask(total_len, num_blocks, block_size, time_token_start_indices, time_token_end_indices, time_token_indices, visual_token_start_pos, visual_token_end_pos, attention_mask, inputs_embeds, prev_blocks_num)

        # print(f'frames:{frames_num}, block_num:{num_blocks}, bsz:{block_size}, prev_blocks_num:{prev_blocks_num}, ratio:{ratio}') 

        return super().forward(
            input_ids=input_ids,
            attention_mask=final_mask, # final_mask
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            labels=labels,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        key_position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        images: Optional[torch.FloatTensor] = None,
        image_sizes: Optional[List[List[int]]] = None,
        return_dict: Optional[bool] = None,
        modalities: Optional[List[str]] = ["image"],
        dpo_forward: Optional[bool] = False,
        cache_position=None,
        time_embedding=None,
        visual_token_start_pos=None,
        visual_token_end_pos=None,
        time_token_start_indices=None,
        frames_num=None,
        time_token_indices=None,
        time_token_end_indices=None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        
        if input_ids is not None and input_ids.size(1) == 1:
            past_key_len = past_key_values[0][0].size(-2)
            key_position_ids = torch.arange(0, past_key_len+1, device=position_ids.device,dtype=torch.long).expand(1, -1)
            if position_ids[0][0] != past_key_len:
                position_ids = torch.tensor([[past_key_len]]).to(device=position_ids.device, dtype=position_ids.dtype)
                key_position_ids = torch.arange(0, past_key_len+1, device=position_ids.device,dtype=torch.long).expand(1, -1)

            return super().forward(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                key_position_ids=key_position_ids,
                past_key_values=past_key_values,
                inputs_embeds=inputs_embeds,
                labels=labels,
                use_cache=use_cache,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )

        if inputs_embeds is None:
            (input_ids, position_ids, attention_mask, past_key_values, inputs_embeds, labels) = self.prepare_inputs_labels_for_multimodal(input_ids, position_ids, attention_mask, past_key_values, labels, images, modalities, image_sizes, time_embedding)

        if self.config.enable_chunk_prefill:

            prefill_mode = self.config.prefill_config['chunk_prefill_mode']
            chunk_size = self.config.prefill_config['chunk_size']
            step_size = self.config.prefill_config['step_size']
            offload = self.config.prefill_config['offload']
          
            if prefill_mode=='streaming':
                return self.forward_streaming(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    key_position_ids=key_position_ids,
                    past_key_values=past_key_values,
                    inputs_embeds=inputs_embeds,
                    labels=labels,
                    use_cache=use_cache,
                    output_attentions=output_attentions,
                    output_hidden_states=output_hidden_states,
                    return_dict=return_dict,
                    cache_position=cache_position,
                    visual_token_start_pos=visual_token_start_pos,
                    visual_token_end_pos=visual_token_end_pos,
                    time_token_start_indices=time_token_start_indices,
                    frames_num=frames_num,
                    time_token_indices=time_token_indices,
                    time_token_end_indices=time_token_end_indices,
                    block_size_chosed=chunk_size,
                    prev_blocks_num=chunk_size - step_size,
                    offload=offload,
                )
            elif prefill_mode=='mask':
                return self.forward_mask(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    past_key_values=past_key_values,
                    inputs_embeds=inputs_embeds,
                    labels=labels,
                    use_cache=use_cache,
                    output_attentions=output_attentions,
                    output_hidden_states=output_hidden_states,
                    return_dict=return_dict,
                    cache_position=cache_position,
                    visual_token_start_pos=visual_token_start_pos,
                    visual_token_end_pos=visual_token_end_pos,
                    time_token_start_indices=time_token_start_indices,
                    frames_num=frames_num,
                    time_token_indices=time_token_indices,
                    time_token_end_indices=time_token_end_indices,
                    block_size_chosed=block_size_chosed,
                    prev_blocks_num=prev_blocks_num,
                )
        else:
            return super().forward(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                inputs_embeds=inputs_embeds,
                labels=labels,
                use_cache=use_cache,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )

        
    @torch.no_grad()
    def generate(
        self,
        inputs: Optional[torch.Tensor] = None,
        images: Optional[torch.Tensor] = None,
        image_sizes: Optional[torch.Tensor] = None,
        modalities: Optional[List[str]] = ["image"],
        time_embedding=None,
        **kwargs,
    ) -> Union[GenerateOutput, torch.LongTensor]:

        position_ids = kwargs.pop("position_ids", None)
        attention_mask = kwargs.pop("attention_mask", None)

        if "inputs_embeds" in kwargs:
            raise NotImplementedError("`inputs_embeds` is not supported")

        if images is not None and images[0].size(0) > 0:
            IMAGE_TOKEN_INDEX = -200
            TOKEN_PERFRAME = 36
            frames_num = images[0].size(0)
            visual_token_start_pos = (inputs == IMAGE_TOKEN_INDEX).nonzero(as_tuple=True)[1].item()
            num_tokens = time_embedding[0].size(0)
            visual_token_end_pos = visual_token_start_pos + num_tokens
            kwargs['visual_token_start_pos'] = visual_token_start_pos
            kwargs['visual_token_end_pos'] = visual_token_end_pos
            # time_token_start_indices = (time_embedding[0] == 1462).nonzero(as_tuple=True)
            time_token_start_indices = (time_embedding[0] == 1462).nonzero(as_tuple=True)[0].cpu().tolist()
            kwargs['time_token_start_indices'] = [idx + visual_token_start_pos for idx in time_token_start_indices]
            # kwargs['time_token_start_indices'] = time_token_start_indices + visual_token_start_pos
            kwargs['frames_num'] = frames_num
            time_token_indices = (time_embedding[0] != 151654).nonzero(as_tuple=True)[0].cpu().tolist()
            kwargs['time_token_indices'] = [idx + visual_token_start_pos for idx in time_token_indices]
            time_token_end_indices = (time_embedding[0] == 25).nonzero(as_tuple=True)[0].cpu().tolist()
            kwargs['time_token_end_indices'] = [idx + visual_token_start_pos + 1 for idx in time_token_end_indices]
            # kwargs['time_token_end_indices'] = time_token_end_indices + visual_token_start_pos

        #print(images[0].shape)
        if images is not None:
            (inputs, position_ids, attention_mask, _, inputs_embeds, _) = self.prepare_inputs_labels_for_multimodal(inputs, position_ids, attention_mask, None, None, images, modalities, image_sizes=image_sizes,time_embedding=time_embedding)
        
        else:
            inputs_embeds = self.get_model().embed_tokens(inputs)
        
        #print(inputs_embeds.shape)
        return super().generate(position_ids=position_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, **kwargs)

    @torch.no_grad()
    def chat(self,
        video_path,
        tokenizer,
        user_prompt,
        chat_history=None,
        return_history=True,
        max_num_frames=512,
        sample_fps=1,
        max_sample_fps=4,
        generation_config={}):

        # prepare text input
        conv = conv_templates["qwen_1_5"].copy()
        if chat_history is None or len(chat_history) == 0:
            user_prompt = f'{DEFAULT_IMAGE_TOKEN}\n{user_prompt}'
        else:
            assert DEFAULT_IMAGE_TOKEN in chat_history[0]['content'], chat_history
            for msg in chat_history:
                conv.append_message(msg['role'], msg['content'])
        
        conv.append_message(conv.roles[0], user_prompt)
        conv.append_message(conv.roles[1], None)
        prompt = conv.get_prompt()
        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(self.model.device)

        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
        keywords = [stop_str]
        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
        generation_config["stopping_criteria"] = [stopping_criteria]

        # prepare video input
        frames, timestamps = load_video(video_path, max_num_frames, fps=sample_fps, max_fps=max_sample_fps)
        print(f'video has loaded, extract {len(frames)} frames.')

        time_stamps=[]
        token_frames_sum=(len(timestamps)+3)//4
        compress_frame = timestamps[::4]
        time_embedding = []
        for time in compress_frame:
            item = f"Time {time}s:"
            time_embedding.append(tokenizer(item).input_ids)
            time_embedding.append([151654]*144)

        time_embedding = [item for sublist in time_embedding for item in sublist]
        time_embedding = torch.tensor(time_embedding, dtype=torch.long).to(self.model.device)
        time_stamps.append(time_embedding)

        video_tensor = self.get_vision_tower().image_processor.preprocess(frames, return_tensors="pt")["pixel_values"].to(self.model.device, dtype=torch.float16)
        
        with torch.inference_mode():
            output_ids = self.generate(input_ids, images=[video_tensor],time_embedding=time_stamps, modalities=["video"], **generation_config)

        outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
        
        if chat_history is None:
            chat_history = []

        chat_history.append({"role":conv.roles[0], "content":user_prompt})
        chat_history.append({"role":conv.roles[1], "content":outputs})
        if return_history:
            return outputs, chat_history
        else:
            return outputs

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
        images = kwargs.pop("images", None)
        image_sizes = kwargs.pop("image_sizes", None)
        visual_token_start_pos = kwargs.get("visual_token_start_pos", None)
        visual_token_end_pos = kwargs.get("visual_token_end_pos", None)
        time_token_start_indices = kwargs.get("time_token_start_indices", None)
        frames_num = kwargs.get("frames_num", None)
        time_token_indices = kwargs.get("time_token_indices", None)
        time_token_end_indices = kwargs.get("time_token_end_indices", None)

        inputs = super().prepare_inputs_for_generation(input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs)

        inputs["visual_token_start_pos"] = visual_token_start_pos
        inputs["visual_token_end_pos"] = visual_token_end_pos
        inputs["time_token_start_indices"] = time_token_start_indices
        inputs["frames_num"] = frames_num
        inputs["time_token_indices"] = time_token_indices
        inputs["time_token_end_indices"] = time_token_end_indices

        if images is not None:
            inputs["images"] = images
        if image_sizes is not None:
            inputs["image_sizes"] = image_sizes
        return inputs


AutoConfig.register("llava_qwen", LlavaQwenConfig)
AutoModelForCausalLM.register(LlavaQwenConfig, LlavaQwenConfig)