llava/model/language_model/llava_qwen2_slowfast.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
from collections import OrderedDict

import transformers
# from transformers import AutoConfig, AutoModelForCausalLM, LlamaConfig, LlamaModel, LlamaForCausalLM
from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.generation.utils import GenerateOutput
from transformers import AutoConfig, AutoModelForCausalLM, Qwen2Config, Qwen2Model, Qwen2ForCausalLM
from transformers.cache_utils import Cache, DynamicCache, StaticCache

from llava.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
from llava.model.llava_arch import LlavaMetaModel, LlavaMetaForCausalLM
from llava.model.language_model.hybrid_decoder_layer import Qwen2DecoderLayer, Qwen2HybridDecoderLayer

class LlavaQwenSlowFastConfig(Qwen2Config):
    model_type = "llava_qwen_slow_fast"


class LlavaQwenSlowFastModel(LlavaMetaModel, Qwen2Model):
    config_class = LlavaQwenSlowFastConfig

    def __init__(self, config: Qwen2Config):
        super(LlavaQwenSlowFastModel, self).__init__(config)
        
        # initialize the cross-attention layers
        self.slow_branch_is_initialized = False
        
        if hasattr(config, "cross_attn_every_n_layers"):
            self.initialize_slow_branch_modules(config)
    
    def initialize_slow_branch_modules(self, args):
        if self.slow_branch_is_initialized:
            return 
        # number of decoder layers
        num_layers = len(self.layers) 
        
        cross_attn_every_n_layers = args.cross_attn_every_n_layers
        cross_attn_gating_type = args.cross_attn_gating_type
        cross_attn_implementation = args.cross_attn_implementation
        cross_attn_max_layer_depth = getattr(args, "cross_attn_max_layer_depth", num_layers)
        cross_attn_min_layer_depth = getattr(args, "cross_attn_min_layer_depth", 0)
        if cross_attn_max_layer_depth is None:
            cross_attn_max_layer_depth = num_layers
        if cross_attn_min_layer_depth is None:
            cross_attn_min_layer_depth = 0
        
        self.config.cross_attn_every_n_layers = cross_attn_every_n_layers
        self.config.cross_attn_implementation = cross_attn_implementation
        self.config.cross_attn_gating_type = cross_attn_gating_type
        self.config.cross_attn_max_layer_depth = cross_attn_max_layer_depth
        self.config.cross_attn_min_layer_depth = cross_attn_min_layer_depth
        
        # set pooling operations
        tile_image_input = getattr(args, "tile_image_input", True) # tile all the image input into a video sequence
        min_fast_frames = getattr(args, "min_fast_frames", 1)   # force to sample at least `min_fast_frames` frames for fast visual tokens
        if min_fast_frames is None:
            min_fast_frames = 1
        
        self.config.tile_image_input = tile_image_input
        self.config.min_fast_frames = min_fast_frames
        
        # generate layer index for the hybrid layer
        hybrid_layer_idx = []
        for i in range(cross_attn_min_layer_depth, cross_attn_max_layer_depth, cross_attn_every_n_layers):
            hybrid_layer_idx.append(i)
        
        # substitute the original decoder layer with hybrid layer
        initialize_kv_from_lm = getattr(args, "initialize_cross_attn_kv_from_lm", False) # whether use LLM's pretrained kv projection to initialize the kv projection weight of cross-attn
        for idx in range(len(self.layers)):
            if idx in hybrid_layer_idx:
                original_decoder_layer = self.layers[idx]
                hybrid_decoder_layer = Qwen2HybridDecoderLayer(self.config, layer_idx=idx, is_hyper_enabled=True, cross_attn_gating_type=cross_attn_gating_type, cross_attn_implementation=cross_attn_implementation)
                _, unexpected_keys = hybrid_decoder_layer.load_state_dict(original_decoder_layer.state_dict(), strict=False) # cause problem when using deepspeed zero3
                if initialize_kv_from_lm and hasattr(hybrid_decoder_layer.self_attn, "cross_attn_kv_proj"):
                    kv_weight = torch.cat([original_decoder_layer.self_attn.k_proj.weight,
                                           original_decoder_layer.self_attn.v_proj.weight], dim=0)
                    kv_bias = torch.cat([original_decoder_layer.self_attn.k_proj.bias,
                                           original_decoder_layer.self_attn.v_proj.bias], dim=0)
                    new_state_dict = OrderedDict()
                    new_state_dict['weight'] = kv_weight
                    new_state_dict['bias'] = kv_bias
                    hybrid_decoder_layer.self_attn.cross_attn_kv_proj.load_state_dict(new_state_dict)
                assert len(unexpected_keys) == 0
                self.layers[idx] = hybrid_decoder_layer
        
        # fast token config
        self.config.fast_token_spatial_stride = args.fast_token_spatial_stride
        self.config.fast_token_temporal_stride = args.fast_token_temporal_stride
        self.config.fast_token_temporal_sampling_stride = args.fast_token_temporal_sampling_stride
        
        self.slow_branch_is_initialized = True

    def forward(
        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,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError(
                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
            )

        if self.gradient_checkpointing and self.training:
            if use_cache:
                # logger.warning_once(
                #     "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                # )
                use_cache = False

        # kept for BC (non `Cache` `past_key_values` inputs)
        return_legacy_cache = False
        if use_cache and not isinstance(past_key_values, Cache):
            return_legacy_cache = True
            if past_key_values is None:
                past_key_values = DynamicCache()
            else:
                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
                # logger.warning_once(
                #     "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and "
                #     "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class "
                #     "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"
                # )

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if cache_position is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(
                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
            )
        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        causal_mask = self._update_causal_mask(
            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
        )

        hidden_states = inputs_embeds

        # create position embeddings to be shared across the decoder layers
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        # decoder layers
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        next_decoder_cache = None

        for decoder_layer in self.layers:
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            if self.gradient_checkpointing and self.training:
                if not isinstance(decoder_layer, Qwen2HybridDecoderLayer):
                    layer_outputs = self._gradient_checkpointing_func(
                        decoder_layer.__call__,
                        hidden_states,
                        causal_mask,
                        position_ids,
                        past_key_values,
                        output_attentions,
                        use_cache,
                        cache_position,
                        position_embeddings,
                    )
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=causal_mask,
                    position_ids=position_ids,
                    past_key_value=past_key_values,
                    output_attentions=output_attentions,
                    use_cache=use_cache,
                    cache_position=cache_position,
                    position_embeddings=position_embeddings,
                )

            hidden_states = layer_outputs[0]

            if use_cache:
                next_decoder_cache = layer_outputs[2 if output_attentions else 1]

            if output_attentions:
                all_self_attns += (layer_outputs[1],)

        hidden_states = self.norm(hidden_states)

        # add hidden states from the last decoder layer
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        next_cache = next_decoder_cache if use_cache else None
        if return_legacy_cache:
            next_cache = next_cache.to_legacy_cache()

        if not return_dict:
            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
        )



class LlavaQwenSlowFastForCausalLM(Qwen2ForCausalLM, LlavaMetaForCausalLM):
    config_class = LlavaQwenSlowFastConfig

    def __init__(self, config):
        Qwen2ForCausalLM.__init__(self, config)
        config.model_type = "llava_qwen_slow_fast"
        config.rope_scaling = None

        self.model = LlavaQwenSlowFastModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_model(self):
        return self.model

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, Qwen2HybridDecoderLayer):
            module.gradient_checkpointing = value
    
    def forward(
        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,
        images: Optional[torch.FloatTensor] = None,
        image_sizes: Optional[List[List[int]]] = None,
        return_dict: Optional[bool] = None,
        modalities: Optional[List[str]] = ["image"],
        cache_position=None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:

        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, image_sizes)

        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"],
        **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:
            (inputs, position_ids, attention_mask, _, inputs_embeds, _) = self.prepare_inputs_labels_for_multimodal(inputs, position_ids, attention_mask, None, None, images, image_sizes=image_sizes)
        else:
            inputs_embeds = self.get_model().embed_tokens(inputs)

        return super().generate(position_ids=position_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, **kwargs)

    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)
        inputs = super().prepare_inputs_for_generation(input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs)
        if images is not None:
            inputs["images"] = images
        if image_sizes is not None:
            inputs["image_sizes"] = image_sizes
        return inputs

    def sample_fast_frames(self,
                      total_frames, 
                      stride,
                      min_frame_number):
    
        all_indices_list = list(range(total_frames))
        
        if total_frames < min_frame_number:
            return all_indices_list
        
        sampled_frames = max(total_frames // stride, min_frame_number)
        stride = total_frames / sampled_frames
        
        fast_indices = [min(int(i * stride), total_frames-1) for i in range(sampled_frames)]
    
        return fast_indices

    def split_slow_fast_tokens(self, 
                               visual_tokens, 
                               temporal_sampling_stride=1,
                               spatial_stride=1,
                               temporal_stride=1):
        # TODO: Min: this function is very messy and can be simplified.
        if isinstance(visual_tokens, torch.Tensor):
            # for all image inputs, only perform spatial pooling 
            b, n, c = visual_tokens.shape
            h = w = int(n**0.5)
            fast_visual_tokens = nn.functional.avg_pool2d(visual_tokens.reshape(b, h, w, c).permute(0, 3, 1, 2),
                                                          kernel_size=spatial_stride, 
                                                          stride=spatial_stride).flatten(2,3).transpose(1,2)            
            return fast_visual_tokens, visual_tokens
        else:
            fast_visual_tokens = []
            for sample_ in visual_tokens:
                t, n, c = sample_.shape
                if t > 1: # is a video
                    T_downsampling_rate = temporal_sampling_stride * temporal_stride
                    
                    if t % T_downsampling_rate != 0:
                        padding_size = (T_downsampling_rate - t % T_downsampling_rate) % T_downsampling_rate
                        # Pad on the first dimension (sequence length) with zeros
                        sample_ = nn.functional.pad(sample_, (0, 0, 0, 0, 0, padding_size))  # (dim_pad_left, dim_pad_right, T_pad_left, T_pad_right)
                    
                    # 1. temporal direct sampling
                    if temporal_sampling_stride > 1:
                        fast_token_indices = self.sample_fast_frames(total_frames=t,
                                                    stride=temporal_sampling_stride,
                                                    min_frame_number=self.config.min_fast_frames)
                    else:
                        fast_token_indices = list(range(t))
                    
                    sample_ = torch.stack([sample_[idx] for idx in fast_token_indices], dim=0)     
                    b, n, c = sample_.shape
                    h = w = int(n**0.5)     
                    sample_ = sample_.reshape(b, h, w, c).permute(0, 3, 1, 2)
                    
                    # 2. temporal average pooling
                    if temporal_stride > 1:
                        if (sample_.shape[0] // temporal_stride) >= self.config.min_fast_frames:
                            sample_ = nn.functional.avg_pool3d(sample_.transpose(0, 1), kernel_size=(temporal_stride, 1, 1)).transpose(0, 1)
                        else:
                            h_, w_ = sample_.shape[-2:]
                            output_frames_num = min(sample_.shape[0], self.config.min_fast_frames)
                            sample_ = nn.functional.adaptive_avg_pool3d(sample_.transpose(0, 1), output_size=(output_frames_num, h_, w_)).transpose(0, 1)
                        
                    # 3. spatial pooling
                    if spatial_stride > 1:
                        sample_ = nn.functional.avg_pool2d(sample_,
                                                        kernel_size=spatial_stride, 
                                                        stride=spatial_stride)
                    sample_ = sample_.flatten(2,3).transpose(1,2)
                
                else:
                    if spatial_stride > 1:
                        h = w = int(n**0.5)     
                        sample_ = sample_.reshape(t, h, w, c).permute(0, 3, 1, 2)
                        sample_ = nn.functional.avg_pool2d(sample_,
                                                        kernel_size=spatial_stride, 
                                                        stride=spatial_stride)
                        sample_ = sample_.flatten(2,3).transpose(1,2)
                        
                fast_visual_tokens.append(sample_.flatten(0, 1).contiguous())
            slow_visual_tokens = [_.flatten(0, 1).contiguous() for _ in visual_tokens]
            
            return fast_visual_tokens, slow_visual_tokens

                
    def prepare_inputs_labels_for_multimodal(
        self, input_ids, position_ids, attention_mask, past_key_values, labels,
        images, image_sizes=None
    ):
        vision_tower = self.get_vision_tower()
        if vision_tower is None or images is None or input_ids.shape[1] == 1:
            # clear the visual tokens if current one is a pure text sample
            if images is None and input_ids.shape[1] > 1:
                for layer in self.get_decoder().layers:
                        if hasattr(layer, "clear_vis_x"):
                            layer.clear_vis_x()
            
            token_types = torch.ones_like(input_ids, dtype=input_ids.dtype, device=input_ids.device)
            for layer in self.get_decoder().layers:
                if hasattr(layer, "condition_vis_x"):
                    layer.media_locations = token_types
            
            return input_ids, position_ids, attention_mask, past_key_values, None, labels

        # handle image input
        images = [image if len(image.shape) == 4 else image.unsqueeze(0) for image in images] # list [ [T, C, H, W], ]
        feature_split_size = [x.shape[0] for x in images]
        all_features, feature_split_size = self.encode_images(torch.cat(images, dim=0), feature_split_size)

        raw_image_features = torch.split(all_features, feature_split_size, dim=0)
        image_features = []
        for sample_feat in raw_image_features:   # initial spatial pooling for all video tokens
            if sample_feat.shape[0] > 1 and self.config.mm_video_pooling_stride > 1:
                b, n, c = sample_feat.shape
                h = w = int(n**0.5)                    
                sample_feat = nn.functional.avg_pool2d(sample_feat.reshape(b, h, w, c).permute(0, 3, 1, 2),
                                                        kernel_size=self.config.mm_video_pooling_stride, 
                                                        stride=self.config.mm_video_pooling_stride).flatten(2,3).transpose(1,2)         
            image_features.append(sample_feat.contiguous())
        del raw_image_features, all_features
        
        ## generate fast and slow tokens
        image_features, slow_image_features = self.split_slow_fast_tokens(image_features,
                                                                        temporal_sampling_stride=self.config.fast_token_temporal_sampling_stride,
                                                                        spatial_stride=self.config.fast_token_spatial_stride,
                                                                        temporal_stride=self.config.fast_token_temporal_stride)
        
        ## set cross-attention states
        if isinstance(slow_image_features, (list, tuple)):
            padded_tensors = torch.nn.utils.rnn.pad_sequence(slow_image_features, batch_first=True)
            cross_attn_mask = torch.ones(padded_tensors.shape[:-1], dtype=torch.bool, device=padded_tensors.device)
            for i, tensor in enumerate(slow_image_features):
                cross_attn_mask[i, len(tensor):] = False  # Mark padded elements as False  
            slow_image_features = padded_tensors
        else:
            cross_attn_mask = torch.ones(slow_image_features.shape[:-1], dtype=torch.bool, device=slow_image_features.device)
        
        # TODO: image start / end is not implemented here to support pretraining.
        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
        new_token_types = [] 
        # NOTE: Min: we need to record the type of tokens so that we can split the tokens in the hybrid decoder layer
        # Token type 1: user's input and system tokens, 2: response text tokens, 3: visual tokens, 4: invalid tokens (padding)
        
        for batch_idx, cur_input_ids in enumerate(input_ids):
            num_images = (cur_input_ids == IMAGE_TOKEN_INDEX).sum()
            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_token_type = torch.full((cur_input_ids.shape[0],), 2, dtype=cur_input_ids[-1].dtype, device=cur_input_ids[-1].device) 
                cur_token_type[labels[batch_idx] == IGNORE_INDEX] = 1 # token with ignore tokens are considered as user input 
                new_token_types.append(cur_token_type)
                
                cur_image_idx += 1
                continue

            image_token_indices = [-1] + torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0].tolist() + [cur_input_ids.shape[0]]
            cur_input_ids_noim = []
            cur_labels = labels[batch_idx]
            cur_labels_noim = []
            cur_token_type_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]])
                
                cur_token = torch.full((cur_labels_noim[-1].shape[0],), 2, dtype=cur_input_ids_noim[-1].dtype, device=cur_input_ids_noim[-1].device) 
                cur_token[cur_labels[image_token_indices[i]+1:image_token_indices[i+1]] == IGNORE_INDEX] = 1 # ingored tokens are considered as user input 
                cur_token_type_noim.append(cur_token) 
                
            split_sizes = [x.shape[0] for x in cur_labels_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 = []
            cur_new_token_type = []

            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])
                cur_new_token_type.append(cur_token_type_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_token_type.append(torch.full((cur_image_features.shape[0],), 3, device=cur_labels.device, dtype=cur_labels.dtype)) # insert image token type

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

            cur_new_input_embeds = torch.cat(cur_new_input_embeds)
            cur_new_labels = torch.cat(cur_new_labels)
            cur_new_token_type = torch.cat(cur_new_token_type) ##

            new_input_embeds.append(cur_new_input_embeds)
            new_labels.append(cur_new_labels)
            new_token_types.append(cur_new_token_type) ##
            
        # 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)
        if tokenizer_model_max_length is not None:
            new_input_embeds = [x[:tokenizer_model_max_length] for x in new_input_embeds]
            new_labels = [x[:tokenizer_model_max_length] for x in new_labels]
            new_token_types = [x[:tokenizer_model_max_length] for x in new_token_types]

        # 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)
        new_token_types_padded = torch.full((batch_size, max_len), 4, dtype=new_labels[0].dtype, device=new_labels[0].device) ## 4 is invalid token type (padding)
        
        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)
                    new_token_types_padded[i, -cur_len:] = new_token_types[i] ## 
            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_token_types_padded[i, :cur_len] = new_token_types[i]

        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

        # token type
        token_types = new_token_types_padded
        # send token type to cross-attn layers
        if _input_ids is not None and _input_ids.shape[-1] == 1:
            pass
        else:
            if slow_image_features is not None:
                for layer in self.get_decoder().layers:
                    if hasattr(layer, "condition_vis_x"):
                        layer.condition_vis_x(slow_image_features, 
                                              cross_attn_mask,
                                              token_type=token_types)
            else:
                for layer in self.get_decoder().layers:
                    if hasattr(layer, "clear_vis_x"):
                        layer.clear_vis_x() 

        
        return None, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels


AutoConfig.register("llava_qwen_slow_fast", LlavaQwenSlowFastConfig)
AutoModelForCausalLM.register(LlavaQwenSlowFastConfig, LlavaQwenSlowFastForCausalLM)