models/modeling_tarsier.py

from dataclasses import dataclass
from typing import List, Optional, Tuple, Union, Dict, Any
import math

import torch.utils.checkpoint
from torch import nn
import torch.nn.functional as F

from transformers import PreTrainedModel, AutoConfig, AutoModel
from transformers.activations import ACT2FN
from transformers.cache_utils import Cache
from transformers.modeling_outputs import ModelOutput
from transformers.utils import logging
from transformers.configuration_utils import PretrainedConfig
from transformers.dynamic_module_utils import get_class_from_dynamic_module
from transformers.models.auto import AutoModel, AutoModelForCausalLM, CONFIG_MAPPING
from transformers.generation import GenerationMixin

from transformers import LlamaForCausalLM, Qwen2ForCausalLM
# from models.modeling_qwen2 import Qwen2ForCausalLM
from models.modeling_qwen2_vl_fast import Qwen2VLForCausalLM
from models.utils import _pad_input, _unpad_input

logger = logging.get_logger(__name__)


class LlavaConfig(PretrainedConfig):

    model_type = "llava"
    is_composition = False

    def __init__(
        self,
        vision_config=None,
        text_config=None,
        ignore_index=-100,
        image_token_index=32000,
        projector_hidden_act="gelu",
        vision_feature_select_strategy="default",
        vision_feature_layer=-2,
        image_newline_idx=32002,
        image_new_idx=32003,
        projection_head="MLP",
        **kwargs,
    ):
        self.ignore_index = ignore_index
        self.image_token_index = image_token_index
        self.projector_hidden_act = projector_hidden_act
        self.vision_feature_select_strategy = vision_feature_select_strategy
        self.vision_feature_layer = vision_feature_layer
        self.image_newline_idx = image_newline_idx
        self.image_new_idx = image_new_idx
        self.projection_head = projection_head

        self.vision_config = vision_config

        if isinstance(self.vision_config, dict):
            vision_config["model_type"] = (
                vision_config["model_type"] if "model_type" in vision_config else "clip_vision_model"
            )
            if 'auto_map' in vision_config:
                repo_id, class_ref = vision_config['auto_map']['AutoConfig'].split("--")
                config_class = get_class_from_dynamic_module(class_ref, repo_id, **kwargs)
                self.vision_config = config_class(**vision_config)
            elif vision_config["model_type"] in CONFIG_MAPPING:
                self.vision_config = CONFIG_MAPPING[vision_config["model_type"]](**vision_config)                
            else:
                raise ValueError(f'vision_config["model_type"] = {vision_config["model_type"]} not supported!')
        
        self.text_config = text_config

        if isinstance(self.text_config, dict):
            text_config["model_type"] = text_config["model_type"] if "model_type" in text_config else "llama"
            if 'auto_map' in text_config:
                repo_id, class_ref = text_config['auto_map']['AutoConfig'].split("--")
                config_class = get_class_from_dynamic_module(class_ref, repo_id, **kwargs)
                self.text_config = config_class(**text_config)
            elif text_config["model_type"] in CONFIG_MAPPING:
                self.text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config)
            else:
                raise ValueError(f'text_config["model_type"] = {text_config["model_type"]} not supported!')
            

        super().__init__(**kwargs)



@dataclass
# Copied from transformers.models.idefics.modeling_idefics.IdeficsCausalLMOutputWithPast with Idefics->Llava
class LlavaCausalLMOutputWithPast(ModelOutput):

    loss: Optional[torch.FloatTensor] = None
    logits: torch.FloatTensor = None
    past_key_values: Optional[List[torch.FloatTensor]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
    position_ids: Optional[torch.LongTensor] = None
    
def add_split_tokens(image_features, image_newline_embed, image_new_embed):
    num_images, num_image_patches, embed_dim = image_features.shape
    num_height_patches, num_width_patches = int(math.sqrt(num_image_patches)), int(math.sqrt(num_image_patches))

    # add image_newline
    image_features = image_features.view(num_images, num_height_patches, num_width_patches, embed_dim)
    image_features = torch.cat([
        image_features,
        image_newline_embed.expand((num_images, num_height_patches, 1, embed_dim))
    ], dim=2)
    num_image_patches += num_height_patches
    image_features = image_features.view(num_images, num_image_patches, embed_dim)

    # add image_new
    image_features = torch.cat([
        image_features,
        image_new_embed.expand((num_images, 1, embed_dim))
    ], dim = 1)

    return image_features


class LlavaMultiModalProjector(nn.Module):
    def __init__(self, config: LlavaConfig):
        super().__init__()
        self.config = config

        self.linear_1 = nn.Linear(config.vision_config.hidden_size, config.text_config.hidden_size, bias=True)
        self.act = ACT2FN[config.projector_hidden_act]
        self.linear_2 = nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size, bias=True)

        image_newline_idx = torch.tensor([config.image_newline_idx], dtype=torch.long)
        image_new_idx = torch.tensor([config.image_new_idx], dtype=torch.long)
        self.register_buffer('image_newline_idx', image_newline_idx, persistent=False)
        self.register_buffer('image_new_idx', image_new_idx, persistent=False)
        

    def forward(self, image_features, input_embeddings):

        selected_image_feature = image_features[self.config.vision_feature_layer]

        if self.config.vision_feature_select_strategy == "default":
            selected_image_feature = selected_image_feature[:, 1:]
        elif self.config.vision_feature_select_strategy == "full":
            selected_image_feature = selected_image_feature
        else:
            raise ValueError(
                f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}"
            )

        hidden_states = self.linear_1(selected_image_feature)
        hidden_states = self.act(hidden_states)
        hidden_states = self.linear_2(hidden_states)

        image_newline_embed = input_embeddings(self.image_newline_idx).squeeze()
        image_new_embed = input_embeddings(self.image_new_idx).squeeze()
        hidden_states = add_split_tokens(hidden_states, image_newline_embed, image_new_embed)
        return hidden_states

class PixelShuffleMultiModalProjector(nn.Module):
    def __init__(self, config: LlavaConfig):
        super().__init__()
        self.config = config

        self.downsample_ratio = 0.5
        vit_hidden_size = config.vision_config.hidden_size
        llm_hidden_size = config.text_config.hidden_size

        self.mlp = nn.Sequential(
            nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),
            nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),
            nn.GELU(),
            nn.Linear(llm_hidden_size, llm_hidden_size)
        )

        image_newline_idx = torch.tensor([config.image_newline_idx], dtype=torch.long)
        image_new_idx = torch.tensor([config.image_new_idx], dtype=torch.long)
        self.register_buffer('image_newline_idx', image_newline_idx, persistent=False)
        self.register_buffer('image_new_idx', image_new_idx, persistent=False)
    
    def forward(self, image_features, input_embeddings):
        selected_image_feature = image_features[self.config.vision_feature_layer]

        if self.config.vision_feature_select_strategy == "default":
            selected_image_feature = selected_image_feature[:, 1:]
        elif self.config.vision_feature_select_strategy == "full":
            selected_image_feature = selected_image_feature
        else:
            raise ValueError(
                f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}"
            )
        
        image_features = self.pixel_shuffle(selected_image_feature)
        hidden_states = self.mlp(image_features)
        
        image_newline_embed = input_embeddings(self.image_newline_idx).squeeze()
        image_new_embed = input_embeddings(self.image_new_idx).squeeze()
        hidden_states = add_split_tokens(hidden_states, image_newline_embed, image_new_embed)

        return hidden_states

    def pixel_shuffle(self, x, scale_factor=0.5):
        if scale_factor == 1:
            return x
        n, wh, c = x.shape
        h, w = int(math.sqrt(wh)), int(math.sqrt(wh))
        x = x.view(n, h, w, c)

        n, w, h, c = x.size()
        # N, W, H, C --> N, W, H * scale, C // scale
        x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
        # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
        x = x.permute(0, 2, 1, 3).contiguous()
        # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
        x = x.view(n, int(h * scale_factor), int(w * scale_factor),
                   int(c / (scale_factor * scale_factor)))
        x = x.permute(0, 2, 1, 3).contiguous()
        x = x.view(x.shape[0], -1, x.shape[-1])
        return x
        

LLAVA_START_DOCSTRING = r"""
    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
    etc.)

    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
    and behavior.

    Parameters:
        config ([`LlavaConfig`] or [`LlavaVisionConfig`]):
            Model configuration class with all the parameters of the model. Initializing with a config file does not
            load the weights associated with the model, only the configuration. Check out the
            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

class TarsierPreTrainedModel(PreTrainedModel):
    config_class = LlavaConfig
    base_model_prefix = "llm"
    supports_gradient_checkpointing = True # TODO: support latest gc
    _skip_keys_device_placement = "past_key_values"
    _supports_flash_attn_2 = True
    _supports_sdpa = False
    _supports_cache_class = True # TODO: support different cache
    _supports_static_cache = True

    def _init_weights(self, module):
        std = (
            self.config.initializer_range
            if hasattr(self.config, "initializer_range")
            else self.config.text_config.initializer_range
        )

        if hasattr(module, "class_embedding"):
            module.class_embedding.data.normal_(mean=0.0, std=std)

        if isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv3d)):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.weight.data.fill_(1.0)
            if module.bias is not None:
                module.bias.data.zero_()
    @property
    def _no_split_modules(self):
        return self.language_model._no_split_modules + self.vision_tower._no_split_modules 


class TarsierForConditionalGeneration(TarsierPreTrainedModel, GenerationMixin):
    def __init__(self, config: LlavaConfig):
        super().__init__(config)
        self.vision_tower = AutoModel.from_config(config.vision_config, trust_remote_code=True)
        if config.text_config.model_type == 'qwen2':
            self.language_model = Qwen2ForCausalLM(config.text_config)
        elif config.text_config.model_type == 'qwen2_vl':
            self.language_model = Qwen2VLForCausalLM(config.text_config)
        elif config.text_config.model_type == 'llama':
            self.language_model = LlamaForCausalLM(config.text_config)
        else:
            raise ValueError(f'{config.text_config.model_type} not supported!')

        if config.projection_head == 'Pixel_Shuffle':
            self.multi_modal_projector = PixelShuffleMultiModalProjector(config)
        elif config.projection_head == 'MLP':
            self.multi_modal_projector = LlavaMultiModalProjector(config)
        elif config.projection_head == 'auto_map':
            repo_id, class_ref = config.auto_map['ProjectionLayer'].split("--")
            model_class = get_class_from_dynamic_module(class_ref, repo_id)
            self.multi_modal_projector = model_class(config)
        elif config.projection_head is None:
            self.multi_modal_projector = lambda x, *args, **kwargs: x

        self.post_init()

    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.language_model.set_input_embeddings(value)

    def get_output_embeddings(self):
        return self.language_model.get_output_embeddings()

    def set_output_embeddings(self, new_embeddings):
        self.language_model.set_output_embeddings(new_embeddings)

    def set_decoder(self, decoder):
        self.language_model.set_decoder(decoder)

    def get_decoder(self):
        return self.language_model.get_decoder()

    def tie_weights(self):
        return self.language_model.tie_weights()

    def resize_token_embeddings(self, new_num_tokens: Optional[int] = None, pad_to_multiple_of=None) -> nn.Embedding:
        model_embeds = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
        # update vocab size
        self.config.text_config.vocab_size = model_embeds.num_embeddings
        return model_embeds

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        pixel_values: torch.FloatTensor = None,
        image_grid_thw: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        labels: Optional[torch.LongTensor] = None,
        num_images: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        use_rmpad: Optional[bool] = False,
        **kwargs,
    ) -> Union[Tuple, LlavaCausalLMOutputWithPast]:
        
        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
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
       
        
        if input_ids is None:
            raise ValueError("You must specify input_ids")
        
        bsz, max_seq_len = input_ids.shape[0], input_ids.shape[1]

        if max_seq_len > 1:
            special_image_mask = input_ids == self.config.image_token_index
            print(f'[{input_ids.device}] num_images: {num_images.tolist()} num_image_tokens: {special_image_mask.sum(-1).tolist()}', flush=True)

        if position_ids is None:
            if 'Qwen2VLForCausalLM' in self.language_model.__class__.__name__:
                position_ids = self.language_model.get_rope_index(input_ids, image_grid_thw, attention_mask) # [bsz, seqlen, 3]
            else:
                position_ids = attention_mask.long().cumsum(-1) - 1 #  # [bsz, seqlen]
                position_ids.masked_fill_(attention_mask == 0, 1)
        
        
        if use_rmpad:
            input_ids, input_ids_indices, cu_seqlens, _ = _unpad_input(input_ids, attention_mask) # [bsz, seqlen] -> [1, seqlen]
            position_ids, _, _, _ = _unpad_input(position_ids, attention_mask)
            input_ids, position_ids = input_ids.unsqueeze(0), position_ids.unsqueeze(0)
        else:
            input_ids_indices, cu_seqlens = None, None

        inputs_embeds = self.get_input_embeddings()(input_ids) # [1, seqlen, dim]
        
        image_features = None
        if pixel_values is not None: # training / first step in generation
            if 'Qwen2VLForCausalLM' in self.language_model.__class__.__name__:
                pixel_values = pixel_values.type(self.vision_tower.get_dtype())
                image_features = self.vision_tower(pixel_values, image_grid_thw)
            else:
                image_outputs = self.vision_tower(pixel_values, output_hidden_states=True)
                image_features = self.multi_modal_projector(
                    image_outputs.hidden_states,
                    self.get_input_embeddings(),
                )

            special_image_mask = input_ids == self.config.image_token_index
            if special_image_mask.sum() > 0:
                image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
                inputs_embeds = inputs_embeds.masked_scatter(
                    special_image_mask.unsqueeze(-1).expand_as(inputs_embeds),
                    image_features
                )
            else:
                inputs_embeds = image_features.sum(dim=(0,1)) * 0. + inputs_embeds

        outputs = self.language_model(
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            use_rmpad=use_rmpad,
            cu_seqlens=cu_seqlens,
        )

        logits = outputs[0]

        loss = None
        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            if use_rmpad:
                labels = labels.view(-1)[input_ids_indices.long()]
                shift_labels = torch.cat((labels[1:], labels.new_ones((1))*-100))
                shift_labels.requires_grad = False
                lbl_seq_lens = (cu_seqlens[1:]-1).long()
                shift_labels[lbl_seq_lens] = -100
                loss = loss_fct(logits.squeeze(0), shift_labels)
            else:
                # Shift so that tokens < n predict n
                shift_logits = logits[..., :-1, :].contiguous()
                shift_labels = labels[..., 1:].contiguous()
                # Flatten the tokens
                shift_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
                shift_labels = shift_labels.view(-1)
                # Enable model parallelism
                shift_labels = shift_labels.to(shift_logits.device)
                loss = loss_fct(shift_logits, shift_labels)
        elif use_rmpad: # 训练的时候，就不 unpad logits 了，节省显存。
            logits = _pad_input(logits.squeeze(0), input_ids_indices, bsz, max_seq_len)

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return LlavaCausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            position_ids=position_ids,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids,
        attention_mask=None,
        position_ids=None,
        past_key_values=None,
        cache_position=None,
        use_cache=True,
        pixel_values=None,
        image_grid_thw=None,
        **kwargs,
    ):
        if past_key_values is not None:
            past_length = past_key_values.get_seq_length()
            input_ids = input_ids[:, past_length:]

        model_inputs = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "past_key_values": past_key_values,
            "use_cache": use_cache,
        }
        if kwargs.get('num_images') is not None:
            model_inputs['num_images'] = kwargs['num_images']

        if cache_position[0] == 0:
            # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
            # Otherwise we need pixel values to be passed to model
            model_inputs["pixel_values"] = pixel_values
            model_inputs["image_grid_thw"] = image_grid_thw
        else:
            model_inputs['position_ids'] = position_ids[:, -1, ...].unsqueeze(1).to(device=input_ids.device) + 1
        return model_inputs


    def _update_model_kwargs_for_generation(
        self,
        outputs: ModelOutput,
        model_kwargs: Dict[str, Any],
        is_encoder_decoder: bool = False,
        num_new_tokens: int = 1,
    ) -> Dict[str, Any]:
        model_kwargs = super()._update_model_kwargs_for_generation(
            outputs=outputs,
            model_kwargs=model_kwargs,
            is_encoder_decoder=is_encoder_decoder,
            num_new_tokens=num_new_tokens,
        )

        if getattr(outputs, "position_ids", None) is not None:
            model_kwargs["position_ids"] = outputs.position_ids

        return model_kwargs