import ast import contextlib import gc import json import os from dataclasses import dataclass from functools import partial from itertools import chain from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.distributed as dist import torch.nn as nn from einops import rearrange from timm.layers import LayerNorm, LayerNorm2d from timm.models.regnet import RegStage from torch.nn import CrossEntropyLoss from transformers import ( AutoConfig, AutoModel, AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, ) from transformers.generation.utils import GenerationMixin from transformers.modeling_utils import ( is_fsdp_enabled, is_local_dist_rank_0, no_init_weights, ) from transformers.models.auto import CONFIG_MAPPING from transformers.utils import ModelOutput from .configuration_hyperclovax import HCXVisionConfig from .image_processing_hyperclovax import select_best_resolution EOT = "<|endofturn|>" IMAGE_LOC = "<|dummy3|>" VIDEO_LOC = "<|_unuse_missing_100270|>" def get_rank(): if dist.is_initialized(): return dist.get_rank() return 0 def get_world_size(): if torch.distributed.is_initialized(): world_size = torch.distributed.get_world_size() else: world_size = 1 return world_size def unpad_image(tensor: torch.Tensor, original_size: Tuple[int, int]) -> torch.Tensor: """Unpads a PyTorch tensor of a padded and resized image. This function removes padding from a tensor image that was previously padded and resized. The padding is removed based on the aspect ratio difference between the original and current image dimensions. Args: tensor: The image tensor, assumed to be in CxHxW format. original_size: The original size of the image as (width, height). Returns: The unpadded image tensor. Examples: >>> import torch >>> # Example 1: Unpadding with height padding >>> padded_tensor = torch.randn(1, 64, 48) # Padded tensor (C=1, H=64, W=48) >>> original_size = (32, 32) # Original size (width=32, height=32) >>> unpadded_tensor = unpad_image(padded_tensor, original_size) >>> unpadded_tensor.shape torch.Size([1, 48, 48]) >>> # Example 2: Unpadding with width padding >>> padded_tensor = torch.randn(1, 48, 64) # Padded tensor (C=1, H=48, W=64) >>> original_size = (32, 32) # Original size (width=32, height=32) >>> unpadded_tensor = unpad_image(padded_tensor, original_size) >>> unpadded_tensor.shape torch.Size([1, 48, 48]) """ original_width, original_height = original_size current_height, current_width = tensor.shape[1:] original_aspect_ratio = original_width / original_height current_aspect_ratio = current_width / current_height if original_aspect_ratio > current_aspect_ratio: 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: 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 def get_anyres_image_grid_shape( image_size: Tuple[int, int], grid_pinpoints: Union[str, List[Tuple[int, int]]], patch_size: int, ) -> Tuple[int, int]: """Calculates the image patch grid shape after any-resolution preprocessing. Selects the optimal resolution from predefined grid pinpoints based on input image dimensions using `select_best_resolution`, then computes the grid layout by dividing the selected resolution by the patch size using integer division. Args: image_size (Tuple[int, int]): Original image dimensions in (width, height) format. grid_pinpoints (Union[str, List[Tuple[int, int]]]): Accepts either: - List of (height, width) resolution tuples - String representation of list (e.g., "[(224, 224), (336, 336)]") patch_size (int): Spatial dimension of square patches for grid division. Returns: Tuple[int, int]: Grid dimensions as (num_patches_width, num_patches_height). Examples: >>> # Basic case with list input >>> get_anyres_image_grid_shape((1000, 800), [(224, 224), (448, 448)], 112) (4, 4) >>> # Basic case with string input >>> get_anyres_image_grid_shape((600, 400), "[(336, 336), (672, 672)]", 112) (6, 6) >>> # Case where resolution is not perfectly divisible by patch_size >>> # select_best_resolution picks (224, 224). 224 // 100 = 2 >>> get_anyres_image_grid_shape((500, 500), [(224, 224)], 100) (2, 2) >>> # Different patch size >>> # select_best_resolution picks (448, 448). 448 // 224 = 2 >>> get_anyres_image_grid_shape((1200, 900), [(448, 448), (224, 224)], 224) (2, 2) Note: String-formatted grid_pinpoints are converted via ast.literal_eval. Invalid formats may raise syntax exceptions. The actual resolution selection depends on the implementation of `select_best_resolution`. The doctests assume `select_best_resolution` picks the *first* resolution provided in `grid_pinpoints`. """ possible_resolutions = grid_pinpoints if isinstance(grid_pinpoints, list) else ast.literal_eval(grid_pinpoints) original_width, original_height = image_size height, width = select_best_resolution((original_height, original_width), possible_resolutions) return width // patch_size, height // patch_size def reshape_and_unpad_image_features( image_feature: torch.Tensor, height: int, width: int, image_size: Tuple[int, int], possible_resolutions: List[Tuple[int, int]], grid_size: int, unpad: bool, image_newline: torch.Tensor, ) -> torch.Tensor: """Reshapes and processes image features with optional unpadding operation. Processes input image features by: 1. Separating base features from spatial features 2. Reshaping spatial features into a 5D tensor (num_patch_height, num_patch_width, height, width, channels) 3. Performing either unpadding operation or simple reshaping based on 'unpad' flag 4. Concatenating processed features with base features Args: image_feature: Input tensor containing image features with shape [1 + num_patches, feature_dim] where the first element is the base feature height: Original image height in pixels width: Original image width in pixels image_size: Target image size as (width, height) tuple possible_resolutions: List of possible [height, width] resolutions for multi-scale processing grid_size: Grid dimension for patch arrangement unpad: Flag to enable unpadding operation image_newline: Special token tensor used as separator when unpadding Returns: torch.Tensor: Processed image features tensor with shape [1 + num_processed_patches, feature_dim] Raises: AssertionError: If base feature dimension doesn't match height*width """ base_image_feature = image_feature[0] image_feature = image_feature[1:] assert ( height * width == base_image_feature.shape[0] ), f"height: {height}, width: {width}, base_image_feature.shape[0]: {base_image_feature.shape[0]}" num_patch_width, num_patch_height = get_anyres_image_grid_shape(image_size, possible_resolutions, grid_size) image_feature = image_feature.view(num_patch_height, num_patch_width, height, width, -1) if unpad: 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_size) image_feature = torch.cat( ( image_feature, 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) image_feature = torch.cat((base_image_feature, image_feature), dim=0) return image_feature def anyres_postprocessing( image_forward_outs: List[torch.FloatTensor], image_sizes: List[List[int]], possible_resolutions: List[Tuple[int, int]], patch_size: int, grid_size: int, image_newline: torch.FloatTensor, num_queries_vis_abstractor: int = -1, unpad: bool = False, ) -> List[torch.FloatTensor]: """Processes 2D visual features into 1D sequences with post-processing steps. Performs AnyRes postprocessing by flattening 2D visual features from grid partitions into 1D sequences, adding newline embeddings at row boundaries for images, and optionally removing padding regions based on original image sizes. For video data, processes each frame's features separately into a single sequence per video and disables unpadding and newline insertion. Args: image_forward_outs (List[torch.FloatTensor]): List of input tensors with shape (number_of_images_in_grid, total_patches, feature_dim) containing visual features. split_sizes (List[int]): A list containing the number of patches for each sample in the batch. The sum of `split_sizes` should equal `image_forward_outs.shape[0]`. image_sizes (List[List[int]]): A list where each element is a list `[width, height]` representing the original dimensions of the corresponding image sample. Used for unpadding. possible_resolutions (List[Tuple[int, int]]): A list of supported resolution tuples `(height, width)` used by `reshape_and_unpad_image_features` for spatial reconstruction, especially during unpadding. patch_size (int): The spatial dimension (height and width) of the square patches the image was divided into. grid_size (int): The spatial dimension (height and width) of the square grid onto which patches are mapped. `grid_size` should be divisible by `patch_size`. image_newline (torch.FloatTensor): A learnable tensor representing the newline embedding, typically with shape (1, feature_dim). Added after each row of image patches when not unpadding. num_queries_vis_abstractor (int, optional): If a visual abstractor with a fixed number of output queries is used instead of grid patching, this specifies the number of queries. Must be a perfect square if > 0. Defaults to -1 (indicating standard grid patching is used). unpad (bool, optional): If `True`, removes padding tokens from image features based on `image_sizes` and `possible_resolutions`. Does not apply to video features. Defaults to False. Returns: List[torch.FloatTensor]: A list of tensors, where each tensor represents the processed 1D sequence of visual features for a single sample from the input batch. The length of the sequence varies depending on processing (unpadding, newlines, video flattening). Raises: AssertionError: If `num_queries_vis_abstractor` is greater than 0 but not a perfect square. """ height = width = grid_size // patch_size if num_queries_vis_abstractor > 0: assert (num_queries_vis_abstractor**0.5).is_integer(), "n_queries must be square number" height = width = int(num_queries_vis_abstractor**0.5) # post-processing (unpad, add newline) new_image_features = [] for image_idx, image_feature in enumerate(image_forward_outs): if image_feature.shape[0] > 1: image_feature = reshape_and_unpad_image_features( image_feature=image_feature, height=height, width=width, image_size=image_sizes[image_idx], possible_resolutions=possible_resolutions, grid_size=grid_size, # Pass grid info if needed by helper unpad=unpad, image_newline=image_newline, ) else: image_feature = image_feature[0] image_feature = torch.cat((image_feature, image_newline[None].to(image_feature.device)), dim=0) new_image_features.append(image_feature) image_features = new_image_features return image_features @dataclass class HCXVisionOutput(ModelOutput): """Output class for vision models, containing various computation results. Args: loss (Optional[torch.FloatTensor], optional): Total cross-entropy loss calculated from logits and labels. loss_per_sample (Optional[torch.FloatTensor], optional): Per-sample loss values for advanced loss processing. logits (torch.FloatTensor): Classification scores (before SoftMax) of shape (batch_size, num_classes). past_key_values (Optional[Tuple[Tuple[torch.FloatTensor]]], optional): Contains precomputed hidden-states that can be used (see `past_key_values` input) to speed up sequential decoding. hidden_states (Optional[Tuple[torch.FloatTensor]], optional): Tuple of torch.FloatTensor (one for the output of the embeddings + one for the output of each layer) of shape (batch_size, sequence_length, hidden_size). Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (Optional[Tuple[torch.FloatTensor]], optional): Tuple of torch.FloatTensor (one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length). Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. """ loss: Optional[torch.FloatTensor] = None loss_per_sample: Optional[torch.FloatTensor] = None logits: torch.FloatTensor = None past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None hidden_states: Optional[Tuple[torch.FloatTensor]] = None attentions: Optional[Tuple[torch.FloatTensor]] = None class HCXVisionForCausalLM(PreTrainedModel, GenerationMixin): """HCX Vision model for causal language modeling with vision-language capabilities. This class combines a vision model with a language model to create a multimodal model capable of processing images or videos and generating text based on the visual inputs. Attributes: config_class: Configuration class for the model. vision_model_name: Name of the vision model component. _no_split_modules: List of modules that should not be split during parallel processing. supports_gradient_checkpointing: Whether the model supports gradient checkpointing. _skip_keys_device_placement: Keys to skip during device placement. """ config_class = HCXVisionConfig vision_model_name = "vision_model" _no_split_modules = ["SiglipEncoderLayer", "LlamaDecoderLayer", "HyperCLOVAXDecoderLayer"] supports_gradient_checkpointing = True _skip_keys_device_placement = "past_key_values" _supports_flash_attn_2 = True _supports_sdpa = True def __init__( self, config: HCXVisionConfig, **kwargs: Optional[Any], ) -> None: """Initialize the HCXVisionForCausalLM model. Args: config: Configuration object for the model containing parameters for both vision and language components. **kwargs: Additional keyword arguments: - use_liger: Whether to use liger kernel for hyperclovax models. - use_fused_ce: Whether to use fused cross-entropy loss. - use_sum_loss: Whether to use sum reduction for loss instead of mean. - is_safetensor_save: Whether to save model using safetensors format. Raises: ValueError: If vision_config is not defined or if text_config is not defined. """ super().__init__(config) # self.config = config # init configs text_config = self._init_text_config(config) vision_config = self._init_vision_config(config) ## possible_resolution should be matched with preprocessor_config.json config.possible_resolutions = self._init_possible_resolutions(config, vision_config) # init models & parameters with no_init_weights(): # weight will be loaded in from_pretrained self.vision_model = AutoModel.from_config(vision_config, trust_remote_code=True) self.mm_projector = self._init_mm_projector(config, text_config, vision_config) self.language_model = AutoModelForCausalLM.from_config(text_config) self.lm_head_vocab_size = getattr(text_config, "padded_vocab_size", text_config.vocab_size) self.language_model.lm_head = nn.Linear(text_config.hidden_size, self.lm_head_vocab_size, bias=False) if config.anyres: self.image_newline = nn.Parameter(torch.empty(text_config.hidden_size, dtype=self.dtype)) # modify configs or model settings if text_config.model_type in ["llama", "hyperclovax", "gpt2"]: self.language_model.gradient_checkpointing_enable() if text_config.model_type == "hyperclovax" and self.use_liger: self.language_model._get_apply_liger_kernel_converter()(model=self.language_model) # update configs self.vision_config = vision_config = self.vision_model.config self.text_config = text_config = self.language_model.config config.update({"vision_config": vision_config}) config.update({"text_config": text_config}) # etc self.use_liger = kwargs.pop("use_liger", False) self.use_fused_ce = kwargs.pop("use_fused_ce", False) self.use_meansum_loss = kwargs.pop("use_meansum_loss", False) self.freeze_before_sampler = kwargs.pop("freeze_before_sampler", False) self.use_turnmeansum_loss = kwargs.pop("use_turnmeansum_loss", False) self.vision_input_chunk_size = kwargs.pop("vision_input_chunk_size", None) self.is_safetensor_save = kwargs.get("is_safetensor_save", True) use_sum_loss = True if kwargs.pop("use_sum_loss", False) else False self.reduction = self._init_reduction_type(use_sum_loss) self.vision_model_use_no_grad = None # forward 시 체크 및 할당 self._backward_compatibility_gradient_checkpointing() # self.post_init() 에 포함되어 있는 gc 가능한지 확인하고 켜주는 함수 def _init_weights(self, module): # copies from https://github.com/kakaobrain/honeybee/blob/main/honeybee/common_layers.py#L55 if ( isinstance(module, nn.Conv2d) # noqa: SIM101 or isinstance(module, nn.Embedding) or isinstance(module, nn.Linear) ): module.weight.data.normal_(mean=0.0, std=0.02) if hasattr(module, "bias") and module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.LayerNorm): module.bias.data.zero_() module.weight.data.fill_(1.0) elif isinstance(module, nn.Parameter): embed_std = 1 / torch.sqrt(torch.tensor(module.size(0), dtype=torch.float)).to(module.dtype) module.data.normal_(mean=0.0, std=embed_std) def _init_reduction_type(self, use_sum_loss): assert not ( self.use_meansum_loss and self.use_turnmeansum_loss ), "use_meansum_loss and use_turnmeansum_loss cannot both be True; only one or neither may be True." if self.use_meansum_loss or self.use_turnmeansum_loss: reduction = "none" elif use_sum_loss: reduction = "sum" else: reduction = "mean" return reduction def _init_vision_config(self, config): vision_model_type = config.vision_config.model_type if vision_model_type in CONFIG_MAPPING: vision_config = CONFIG_MAPPING[vision_model_type](**config.vision_config.to_dict()) vision_config.auto_map = {} else: if config.vision_model_name_or_path is not None: vision_config = AutoConfig.from_pretrained(config.vision_model_name_or_path, trust_remote_code=True) elif config.vision_config._name_or_path is not None: vision_config = AutoConfig.from_pretrained(config.vision_config._name_or_path, trust_remote_code=True) else: raise ValueError("vision_config is not defined") vision_config.anyres = config.anyres vision_config.max_num_grids = config.max_num_grids return vision_config def _init_text_config(self, config): if hasattr(config, "text_config") and config.text_config is not None: model_type = config.text_config.model_type text_config = CONFIG_MAPPING[model_type](**config.text_config.to_dict()) else: raise ValueError("text_config is not defined") text_config._attn_implementation = config._attn_implementation if text_config.model_type != "hyperclovax": text_config.logits_scaling = 1.0 return text_config def _init_possible_resolutions(self, config, vision_config): """possible_resolution should be matched with preprocessor_config.json""" if not getattr(config, "possible_resolutions", []): possible_resolutions = [] if config.anyres: assert config.max_num_grids > 0 for i in range(1, config.max_num_grids + 1): for j in range(1, config.max_num_grids + 1): if i == 1 and j == 1 and not config.use_1x1_grid: continue if i * j <= config.max_num_grids: possible_resolutions.append([i, j]) possible_resolutions = [ [ys * vision_config.image_size, xs * vision_config.image_size] for ys, xs in possible_resolutions ] return possible_resolutions else: return config.possible_resolutions def _init_mm_projector(self, config, text_config, vision_config): input_hidden_size = vision_config.hidden_size if config.mm_projector_type == "linear": mm_projector = nn.Linear(input_hidden_size, text_config.hidden_size) mm_projector.dtype = next(mm_projector.parameters()).dtype elif config.mm_projector_type == "cabstractor": mm_projector = HCXVisionCAbstractor( num_queries=config.num_queries_vis_abstractor_image, num_input_tokens=(vision_config.image_size // vision_config.patch_size) ** 2, encoder_hidden_size=input_hidden_size, hidden_size=input_hidden_size, output_hidden_size=text_config.hidden_size, pos_emb=config.proj_pos_emb, prenorm=config.proj_prenorm, ) else: mm_projector = HCXVisionMlp( config.mm_projector_type, input_hidden_size, hidden_features=input_hidden_size, # TODO: llava 처럼 hidden_size 를 input_hidden_size 가 아니라 LLM embedding size 로 바꿔주기 out_features=self.text_config.hidden_size, ) return mm_projector def forward( self, input_ids: Optional[torch.LongTensor] = None, pixel_values_images: Optional[List[List[torch.FloatTensor]]] = None, image_sizes_images: Optional[List[List[Tuple[int, int]]]] = None, pixel_values_videos: Optional[List[List[torch.FloatTensor]]] = None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, attention_mask: Optional[torch.FloatTensor] = None, position_ids: Optional[torch.LongTensor] = 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, **kwargs, ) -> Union[Tuple, HCXVisionOutput]: """Forward pass of the model. This method processes the input tokens and images, combines them into a unified representation, and generates text output based on the inputs. Args: input_ids: Input token IDs. In positions where images are inputted, the value is replaced by "<|dummy3|>" pixel_values: List of lists of 4D tensors for images. Each outer list corresponds to a batch and contains inner lists of image tensors. past_key_values: Pre-computed key and value states of the attention layers for faster inference. attention_mask: Mask to avoid performing attention on padding token indices. inputs_embeds: Input embeddings. If provided, input_ids will not be used. labels: Labels for computing the language modeling loss. use_cache: Whether to use past key/values for faster inference. output_attentions: Whether to return attention weights of each layer. output_hidden_states: Whether to return hidden states of each layer. return_dict: Whether to return a ModelOutput instead of a tuple. image_sizes: List of lists representing image dimensions (width, height). vision_query_lengths: List of lists containing lengths when each image is converted into visual tokens. non_vision_query_lengths: List of lengths of text tokens (excluding visual tokens) for each sample. img_start_ids_list: List of lists containing indices of img_start_id tokens for each sample. num_queries_vis_abstractors: List of lists containing number of visual tokens for each image grid.\ For video frames, this is the number of visual tokens for the fast part. num_queries_vis_abstractors_slow: List of lists containing number of visual tokens for the slow part when applying the slowfast algorithm to video frames. first_last_frames_slows: List of booleans indicating whether the slowfast algorithm is applied to the first or last frames of the video. is_video_list: List of booleans indicating which inputs are videos. **kwargs: Additional keyword arguments. Returns: If return_dict=True, returns an HCXVisionOutput object containing: - loss: Language modeling loss if labels are provided, otherwise None. - loss_per_sample: Per-sample loss if labels are provided, otherwise None. - logits: Prediction scores of the language modeling head. - past_key_values: Past key/values for faster inference if use_cache=True. - hidden_states: Hidden states of all layers if output_hidden_states=True. - attentions: Attention weights of all layers if output_attentions=True. If return_dict=False, returns a tuple containing the above items except loss_per_sample. """ output_attentions = ( output_attentions if output_attentions is not None else self.config.vision_config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.vision_config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict if inputs_embeds is None and past_key_values is None: if pixel_values_images is not None or pixel_values_videos is not None: inputs_embeds = self.extract_inputs_embeds( input_ids=input_ids, pixel_values_images=pixel_values_images, image_sizes_images=image_sizes_images, pixel_values_videos=pixel_values_videos, ) else: inputs_embeds = self.get_input_embeddings()(input_ids) if inputs_embeds is not None: input_ids = None ################################ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.language_model.base_model( input_ids=input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) hidden_states = outputs[0] hidden_states = hidden_states * self.text_config.logits_scaling loss = None loss_per_sample = None logits = self.language_model.lm_head(hidden_states) if labels is not None: # Shift so that tokens < n predict n shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() # Flatten the tokens loss_fct = CrossEntropyLoss(reduction="none") # ignore IGNORE_INDEX(-100) shift_logits = shift_logits.view(-1, self.lm_head_vocab_size) shift_labels = shift_labels.view(-1) # Enable model/pipeline parallelism shift_labels = shift_labels.to(shift_logits.device) loss = loss_fct(shift_logits, shift_labels) if get_rank() == 0: loss_per_sample = loss.view(logits.shape[0], -1).sum(axis=1) / ( shift_labels.view(logits.shape[0], -1) != self.config.ignore_index ).sum(axis=1) loss = loss[shift_labels != self.config.ignore_index].mean() if not return_dict: output = (logits,) + outputs[1:] return (loss,) + output if loss is not None else output return HCXVisionOutput( loss=loss, loss_per_sample=loss_per_sample, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.get_input_embeddings def get_input_embeddings(self): return self.language_model.get_input_embeddings() # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.set_input_embeddings def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.get_output_embeddings def get_output_embeddings(self): return self.language_model.get_output_embeddings() # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.set_output_embeddings def set_output_embeddings(self, new_embeddings): self.language_model.set_output_embeddings(new_embeddings) # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.set_decoder def set_decoder(self, decoder): self.language_model.set_decoder(decoder) # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.get_decoder def get_decoder(self): return self.language_model.get_decoder() # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.tie_weights def tie_weights(self): return self.language_model.tie_weights() # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.resize_token_embeddings 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) self.config.text_config.vocab_size = model_embeds.num_embeddings self.vocab_size = model_embeds.num_embeddings return model_embeds def extract_inputs_embeds( self, input_ids: Optional[torch.LongTensor] = None, pixel_values_images: Optional[List[List[torch.FloatTensor]]] = None, image_sizes_images: Optional[List[List[Tuple[int, int]]]] = None, pixel_values_videos: Optional[List[List[torch.FloatTensor]]] = None, ): """Extract input embeddings by processing text tokens and visual features. This method processes the input tokens and image features, extracts the visual features using the vision model, and combines them with the text token embeddings to create a unified input representation for the language model. Args: input_ids: Input token IDs with img_start_id markers for image positions. pixel_values: List of lists of image tensors. past_key_values: Pre-computed key and value states for faster inference. image_sizes: List of lists of image dimensions (width, height). vision_query_lengths: List of lists of lengths when each image is converted to visual tokens. non_vision_query_lengths: List of lengths of text tokens (excluding visual tokens) for each sample. img_start_ids_list: List of lists containing indices of img_start_id tokens for each sample. first_last_frames_slows: List of booleans indicating whether the slowfast algorithm is applied to the first or last frames of the video. is_videos: List of booleans indicating which inputs are videos. Returns: Combined embeddings of text tokens and visual features. """ # for convert back to List of List format len_pixel_values_images = [len(pixel_value) for pixel_value in pixel_values_images] if pixel_values_images else [] len_pixel_values_videos = [len(pixel_value) for pixel_value in pixel_values_videos] if pixel_values_videos else [] if sum(len_pixel_values_images) + sum(len_pixel_values_videos) == 0: return None inputs_embeds = self.get_input_embeddings()(input_ids) if sum(len_pixel_values_images) > 0: image_features_batch = self.forward_images( pixel_values_images, image_sizes_images, len_pixel_values_images ) for i, image_features in enumerate(image_features_batch): if len(image_features) > 0: image_token_indices = (input_ids[i] == self.config.image_token_id).nonzero().squeeze() inputs_embeds[i][image_token_indices] = torch.cat(image_features).to(inputs_embeds.dtype) if sum(len_pixel_values_videos) > 0: video_features_batch = self.forward_videos(pixel_values_videos, len_pixel_values_videos) for i, video_features in enumerate(video_features_batch): if len(video_features) > 0: video_token_indices = (input_ids[i] == self.config.video_token_id).nonzero().squeeze() inputs_embeds[i][video_token_indices] = torch.cat(video_features).to(inputs_embeds.dtype) return inputs_embeds def forward_images( self, pixel_values_images: List[List[torch.FloatTensor]], image_sizes_images: List[List[Tuple[int, int]]], len_pixel_values_images: List[int], ) -> List[List[torch.Tensor]]: if sum(len_pixel_values_images) == 0: return None concat_pixel_values_images = torch.cat(list(chain(*pixel_values_images)), dim=0) visual_token_idx = 0 if "siglip" in self.vision_config.model_type else 1 context_vision_model = torch.no_grad() if self.vision_model_use_no_grad else contextlib.nullcontext() with context_vision_model: if self.config.use_nth_layer == -1: # Replace post_layernorm of the last layer with Identity self.vision_model.vision_model.post_layernorm = nn.Identity() image_forward_outs = self.vision_model(concat_pixel_values_images) image_forward_outs = image_forward_outs.last_hidden_state[:, visual_token_idx:] else: image_forward_outs = self.vision_model(concat_pixel_values_images, output_hidden_states=True) image_forward_outs = image_forward_outs.hidden_states[self.config.use_nth_layer][:, visual_token_idx:] image_forward_outs = image_forward_outs.to(dtype=self.mm_projector.dtype) image_forward_outs = self.mm_projector(image_forward_outs) # b (h w) d # feature 를 분할. e.g. torch.Size([18, 81, 3072]) -> [torch.Size([9, 81, 3072]), torch.Size([9, 81, 3072])] split_sizes = [pixel_value.shape[0] for pixel_value in chain(*pixel_values_images)] image_forward_outs = torch.split(image_forward_outs, split_sizes, dim=0) # newline 붙여주기 (anyres postprocessing) image_features = anyres_postprocessing( image_forward_outs=image_forward_outs, image_sizes=[image_size for image_sizes in image_sizes_images for image_size in image_sizes], num_queries_vis_abstractor=self.config.num_queries_vis_abstractor_image, unpad=self.config.unpad, patch_size=self.vision_config.patch_size, grid_size=self.vision_config.image_size, image_newline=self.image_newline, possible_resolutions=self.config.possible_resolutions, ) # 원래 pixel_values_images 형태로 복원 image_features = [ image_features[sum(len_pixel_values_images[:i]) : sum(len_pixel_values_images[: i + 1])] for i in range(len(len_pixel_values_images)) ] return image_features def forward_videos( self, pixel_values_videos: List[List[torch.FloatTensor]], len_pixel_values_videos: List[int], ) -> List[torch.Tensor]: len_video_grids = sum(len_pixel_values_videos) if len_video_grids == 0: return None # Run Vision Model concat_pixel_values_videos = torch.cat(list(chain(*pixel_values_videos)), dim=0) visual_token_idx = 0 if "siglip" in self.vision_config.model_type else 1 context_vision_model = torch.no_grad() if self.vision_model_use_no_grad else contextlib.nullcontext() with context_vision_model: if self.config.use_nth_layer == -1: # Replace post_layernorm of the last layer with Identity self.vision_model.vision_model.post_layernorm = nn.Identity() video_forward_outs = self.vision_model(concat_pixel_values_videos) video_forward_outs = video_forward_outs.last_hidden_state[:, visual_token_idx:] else: video_forward_outs = self.vision_model(concat_pixel_values_videos, output_hidden_states=True) video_forward_outs = video_forward_outs.hidden_states[self.config.use_nth_layer][:, visual_token_idx:] video_forward_outs = video_forward_outs.to(dtype=self.mm_projector.dtype) # Run MM-Projector # len(num_grids) == len(num_queries_vis_abstractors) + 1 grid_idx = 0 num_grids = [grid_idx] # e.g. [0, 9, 18, 19, 27, 28, 36, 37, 45, 46, 54, 55, 56] num_queries_vis_abstractors = [] # e.g. [81, 81, 81, 9, 81, 9, 81, 9, 81, 9, 81, 9] len_total_frames = video_forward_outs.shape[0] if self.config.first_last_frames_slow: # TODO: 동작 확인 안 했음. 해야 함. # slowfast (first_last_frames_slow) assert len_total_frames != 0 if len_total_frames <= 2: num_queries_vis_abstractors.append(self.config.num_queries_vis_abstractor_video_slow) grid_idx += len_total_frames num_grids.append(grid_idx) else: num_queries_vis_abstractors.append(self.config.num_queries_vis_abstractor_video_slow) grid_idx += 1 num_grids.append(grid_idx) num_queries_vis_abstractors.append(self.config.num_queries_vis_abstractor_video_fast) grid_idx += len_total_frames - 2 num_grids.append(grid_idx) num_queries_vis_abstractors.append(self.config.num_queries_vis_abstractor_video_slow) grid_idx += 1 num_grids.append(grid_idx) else: # slowfast for pixel_values_frames in pixel_values_videos: for pixel_values_frame in pixel_values_frames: if len(pixel_values_frame) > 0: num_queries_vis_abstractors.append(self.config.num_queries_vis_abstractor_video_slow) grid_idx += 1 num_grids.append(grid_idx) num_queries_vis_abstractors.append(self.config.num_queries_vis_abstractor_video_fast) grid_idx = grid_idx + len(pixel_values_frame) - 1 num_grids.append(grid_idx) video_forward_outs = self.mm_projector(video_forward_outs, num_queries_vis_abstractors, num_grids) # video_group 별로 concat 처리. # 예를 들어, 3x3 grid 를 사용했을 경우, 총 9개의 feature 가 모일 때까지, grouped_features 에 리스트를 모아주고, concat 처리. video_features = [] # what we want to return target_features = [] target_group_size = 0 group_counter = 0 video_groups = [ len(frame) for frames in pixel_values_videos for frame in frames ] # for concat video features after projector for forward_out in video_forward_outs: target_group_size += len(forward_out) target_features.append(forward_out.flatten(0, 1)) video_group_size = video_groups[group_counter] if video_group_size == target_group_size: video_features.append(torch.cat(target_features, dim=0)) target_features = [] group_counter += 1 target_group_size = 0 elif video_group_size < target_group_size: raise RuntimeError(f"video_group_size < target_group_size!! [{video_group_size} < {target_group_size}]") assert len(target_features) == 0, f"target_features is not empty!! {target_features}" assert len(video_groups) == len(video_features) # 원래 pixel_values_videos 형태로 복원 video_features = [ video_features[sum(len_pixel_values_videos[:i]) : sum(len_pixel_values_videos[: i + 1])] for i in range(len(len_pixel_values_videos)) ] return video_features @torch.no_grad() def generate( self, input_ids: Optional[torch.LongTensor] = None, pixel_values_images: Optional[List[List[torch.FloatTensor]]] = None, image_sizes_images: Optional[List[List[Tuple[int, int]]]] = None, pixel_values_videos: Optional[List[List[torch.FloatTensor]]] = None, pad_token_id: Optional[int] = None, eos_token_id: Optional[int] = None, bad_words_ids: Optional[List[List[int]]] = None, max_length: int = 196, min_length: int = 2, do_sample: bool = True, num_beams: int = 1, top_p: float = 0.6, top_k: int = 0, temperature: float = 0.5, repetition_penalty: float = 1.0, length_penalty: int = 1, use_cache: bool = True, verbose: bool = False, **kwargs, ) -> torch.LongTensor: """Generate text based on input tokens and images. This method generates text based on the provided input tokens and images using beam search and/or sampling strategies. Args: input_ids: Input token IDs with img_start_id markers for image positions. pixel_values: List of lists of image tensors. image_sizes: List of lists of image dimensions (width, height). vision_query_lengths: List of lists of lengths when each image is converted to visual tokens. non_vision_query_lengths: List of lengths of text tokens (excluding visual tokens) for each sample. num_queries_vis_abstractors: List of lists containing number of visual tokens for each image grid. num_queries_vis_abstractors_slow: List of lists containing number of visual tokens for the slow part when applying the slowfast algorithm to video frames. first_last_frames_slows: List of booleans indicating whether the slowfast algorithm is applied to the first or last frames of the video. is_videos: List of booleans indicating which inputs are videos. img_start_ids_list: List of lists containing indices of img_start_id tokens for each sample. pad_token_id: Token ID used for padding. eos_token_id: Token ID used to signal the end of a sequence. bad_words_ids: List of token ID sequences that should not be generated. max_length: Maximum length of the sequence to be generated (input length + max_new_tokens). min_length: Minimum length of the sequence to be generated (input length + min_new_tokens). do_sample: Whether to use sampling for generation (otherwise uses greedy decoding). num_beams: Number of beams for beam search. 1 means no beam search. top_p: Nucleus sampling parameter. Tokens with cumulative probability > top_p are kept. top_k: Number of highest probability tokens to keep for top-k-filtering. temperature: Value used to modulate the next token probabilities. repetition_penalty: Penalty applied to tokens that have already appeared in the sequence. length_penalty: Exponential penalty applied to sequence length. use_cache: Whether to use past key/values for faster inference. **kwargs: Additional keyword arguments. Returns: Generated token IDs. """ # inputs_embeds: torch.bfloat16 : [batchsize, variable(visual token, text token, system prompt 모두 포함)] if pad_token_id is None: pad_token_id = self.tokenizer.pad_token_id if eos_token_id is None: eos_token_id = self.tokenizer.encode("<|endofturn|>")[0] if bad_words_ids is None: bad_words_ids = [ [ self.config.text_config.bos_token_id, ], [ self.config.text_config.eos_token_id, ], ] if (pixel_values_images is None or all(len(pixel_values) == 0 for pixel_values in pixel_values_images)) and ( pixel_values_videos is None or all(len(pixel_values) == 0 for pixel_values in pixel_values_videos) ): return self.language_model.generate( input_ids, pad_token_id=pad_token_id, eos_token_id=eos_token_id, bad_words_ids=bad_words_ids, **kwargs ) inputs_embeds = self.extract_inputs_embeds( input_ids=input_ids, pixel_values_images=pixel_values_images, image_sizes_images=image_sizes_images, pixel_values_videos=pixel_values_videos, ) inputs_embeds = inputs_embeds.to(device=self.language_model.device, dtype=self.language_model.dtype) # pred : torch.int64 : [batchsize, generated token_length] pred = self.language_model.generate( inputs_embeds=inputs_embeds, pad_token_id=pad_token_id, eos_token_id=eos_token_id, bad_words_ids=bad_words_ids, max_new_tokens=max_length, min_length=min_length, num_beams=num_beams, do_sample=(False if temperature == 0.0 else do_sample), # set do_sample=False if invalid temperature top_k=top_k, top_p=top_p, temperature=temperature, repetition_penalty=repetition_penalty, length_penalty=length_penalty, early_stopping=(False if num_beams <= 1 else True), # set early_stopping=False when not beam_search use_cache=use_cache, ) if verbose: llm_query = self.tokenizer.batch_decode( [ [token_id for token_id in input_ids_row if token_id != self.tokenizer.pad_token_id] for input_ids_row in input_ids.detach().cpu().tolist() ], skip_special_tokens=False, )[0] llm_pred = self.tokenizer.batch_decode( [ [token_id for token_id in pred_row if token_id != self.tokenizer.pad_token_id] for pred_row in pred.detach().cpu().tolist() ], skip_special_tokens=False, )[0] print(f"# [info] llm_query: {llm_query}") print(f"# [info] llm_pred: {llm_pred}") return pred def to_vision_model_device(self, input_tensor: Union[torch.Tensor, List]) -> Union[torch.Tensor, List]: """Move input tensors to the vision model's device. This method recursively moves input tensors or lists of tensors to the vision model's device. Args: input_tensor: Input tensor or list of tensors to be moved to the vision model's device. Returns: The input tensor or list of tensors moved to the vision model's device. Raises: TypeError: If the input is neither a tensor nor a list. """ if isinstance(input_tensor, list): return [self.to_vision_model_device(item) for item in input_tensor] elif isinstance(input_tensor, torch.Tensor): return input_tensor.to(self.vision_model.device) else: raise TypeError("Unsupported data type. Only tensors and lists are allowed.") def prepare_inputs_for_generation( self, input_ids: torch.LongTensor, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, attention_mask: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, **kwargs, ) -> Dict[str, Any]: """Prepare inputs for the generation algorithm. This method prepares the input for each generation step based on the model's needs. Args: input_ids: Input token IDs. past_key_values: Pre-computed key and value states for faster inference. attention_mask: Mask to avoid performing attention on padding token indices. inputs_embeds: Input embeddings. If provided, input_ids will not be used. **kwargs: Additional keyword arguments. Returns: Dictionary containing the prepared inputs for the model. """ input_ids = kwargs.get("decoder_input_ids", input_ids) if past_key_values: input_ids = input_ids[:, -1:] # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: model_inputs = {"input_ids": input_ids} model_inputs.update( { "past_key_values": past_key_values, "use_cache": kwargs.get("use_cache"), "attention_mask": attention_mask, "pixel_values": kwargs.get("pixel_values", None), } ) return model_inputs @classmethod def from_config(cls, config, vision_model_name_or_path): return cls(config, vision_model_name_or_path) @classmethod def from_pretrained( cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None, *model_args, **kwargs, ) -> "HCXVisionForCausalLM": assert pretrained_model_name_or_path is not None save_only_vision = kwargs.pop("save_only_vision") if "save_only_vision" in kwargs else False save_only_qformer = kwargs.pop("save_only_qformer") if "save_only_qformer" in kwargs else False save_shard_size = kwargs.pop("save_shard_size") if "save_shard_size" in kwargs else "5GB" if pretrained_model_name_or_path is not None: # when evaluate or load instruction tunned model model: HCXVisionForCausalLM = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) model.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path) image_token_id = model.tokenizer.encode(IMAGE_LOC, add_special_tokens=False) assert ( len(image_token_id) == 1 ), f'"<|dummy3|>" was not encoded into a single special token. Encoding result: {image_token_id}' model.config.image_token_id = image_token_id[0] video_token_id = model.tokenizer.encode(VIDEO_LOC, add_special_tokens=False) assert ( len(video_token_id) == 1 ), f'"<|_unuse_missing_100270|>" was not encoded into a single special token. Encoding result: {video_token_id}' model.config.video_token_id = video_token_id[0] model.save_only_vision = save_only_vision model.save_only_qformer = save_only_qformer model.save_shard_size = save_shard_size return model def get_language_model(self): return self.language_model.base_model def get_vision_model(self): return self.vision_model def save_pretrained( self, save_directory: Union[str, os.PathLike], *args, **kwargs, ): state_dict = kwargs["state_dict"] if "state_dict" in kwargs else self.state_dict() partial_state_dict = self.get_pretrained_state_dict( state_dict, save_directory, ) kwargs["state_dict"] = partial_state_dict kwargs["safe_serialization"] = self.is_safetensor_save kwargs.setdefault("max_shard_size", self.save_shard_size) super().save_pretrained(save_directory, *args, **kwargs) def get_pretrained_state_dict(self, state_dict, save_dir): vision_key = "vision_model." llm_keys = ["language_model."] head_key = "lm_head." for key in list(state_dict.keys()): if self.save_only_vision: for llm_key in llm_keys: if llm_key in key: state_dict.pop(key) if key.startswith(head_key): state_dict.pop(key) elif self.save_only_qformer: if f"{vision_key}" in key: state_dict.pop(key) return state_dict class HCXVisionMlp(nn.Module): def __init__( self, mm_projector_type, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, ): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features self.mm_projector_type = mm_projector_type if self.mm_projector_type == "mlp": self.fc1 = nn.Linear(in_features, hidden_features) self.act = act_layer() self.fc2 = nn.Linear(hidden_features, out_features) elif self.mm_projector_type == "inverted_mlp": self.fc1 = nn.Linear(in_features, 2 * hidden_features) self.act = act_layer() self.fc2 = nn.Linear(2 * hidden_features, out_features) else: raise NotImplementedError("{} is not implemented".format(self.mm_projector_type)) def forward(self, x): x = self.fc1(x) x = self.act(x) x = self.fc2(x) return x class HCXVisionCAbstractor(nn.Module): """ This module is based on C-Abstractor, whose license is under apache-2.0. You can check the original code at https://github.com/khanrc/honeybee/blob/main/honeybee/projectors/projectors.py and we made necessary modifications. """ def __init__( self, num_queries: int, num_input_tokens: int, encoder_hidden_size: int, hidden_size: int, output_hidden_size: int, pos_emb: bool = True, prenorm: bool = False, ): super().__init__() self.num_input_tokens = num_input_tokens self.output_hidden_size = output_hidden_size # Positional embedding if pos_emb: self.pos_emb = torch.nn.Parameter(torch.zeros(1, num_input_tokens, encoder_hidden_size)) self.pos_emb.data.normal_(mean=0.0, std=0.02) else: self.pos_emb = None # (Optional) Pre-normalization layer if prenorm: self.prenorm = LayerNorm(encoder_hidden_size) else: self.prenorm = None self.build_net(num_queries, encoder_hidden_size, hidden_size, output_hidden_size) self.dtype = next(self.parameters()).dtype def forward( self, x: torch.Tensor, num_queries_vis_abstractors: Optional[List[List[int]]] = None, num_grids: Optional[List[int]] = None, ) -> torch.Tensor: """ Args: x: (B, L, encoder_hidden_size) tensor from the visual backbone (e.g. CLIP visual encoder), including cls token. """ if self.prenorm is not None: x = self.prenorm(x) if self.pos_emb is not None: x = x + self.pos_emb x = self._forward( x, num_queries_vis_abstractors=num_queries_vis_abstractors, num_grids=num_grids, ) # (B, L, output_hidden_size) return x def _forward( self, x: torch.Tensor, num_queries_vis_abstractors: Optional[List[List[int]]] = None, num_grids: Optional[List[int]] = None, ) -> torch.Tensor: # x: [B, L, dim] B, L, dim = x.shape hw = int(L**0.5) x = rearrange(x, "b (h w) d -> b d h w", h=hw, w=hw) if num_queries_vis_abstractors is not None: assert num_grids is not None return self._forward_adaptive_num_query(x, num_queries_vis_abstractors, num_grids) x = self.net(x) x = rearrange(x, "b d h w -> b (h w) d") x = self.readout(x) return x def _forward_adaptive_num_query( self, x: torch.Tensor, num_queries_vis_abstractors: Optional[List[List[int]]] = None, num_grids: Optional[List[int]] = None, ) -> List[torch.Tensor]: # self.net is consisted by 3 layers (s1, sampler, s2) assert len(self.net) == 3 x = self.net[0](x) # s1 new_x = [] for i, num_queries in enumerate(num_queries_vis_abstractors): hw = int(num_queries**0.5) sampler = nn.AdaptiveAvgPool2d((hw, hw)) out = sampler(x[num_grids[i] : num_grids[i + 1], :]) out = self.net[2](out) # s2 out = rearrange(out, "b d h w -> b (h w) d") out = self.readout(out) new_x.append(out) return new_x def build_net( self, n_queries: int, encoder_hidden_size: int, hidden_size: int, output_hidden_size: int, depth: int = 3, mlp_depth: int = 2, ): assert (n_queries**0.5).is_integer(), f"n_queries must be square number. n_queries: {n_queries}" hw = int(n_queries**0.5) # RegBlock = ResBlock + SE RegBlock = partial( RegStage, stride=1, dilation=1, act_layer=nn.SiLU, norm_layer=LayerNorm2d, ) s1 = RegBlock( depth, encoder_hidden_size, hidden_size, ) sampler = nn.AdaptiveAvgPool2d((hw, hw)) s2 = RegBlock( depth, hidden_size, hidden_size, ) self.net = nn.Sequential(s1, sampler, s2) self.readout = self.build_mlp(mlp_depth, hidden_size, output_hidden_size) def build_mlp( self, depth: int, hidden_size: int, output_hidden_size: int, ): layers = [nn.Linear(hidden_size, output_hidden_size)] for _ in range(1, depth): layers.append(nn.SiLU()) layers.append(nn.Linear(output_hidden_size, output_hidden_size)) return nn.Sequential(*layers)