o }oiT.@sddlmZmZddlZddZddZdeded efd d Zd d Z dddZ de de de de deee e fde d e fddZ ddZ dS))ListTupleNcCs*||}||}||}||r|SdS)zHGet image sequence length given image size, patch size, and class token.r)img_himg_w patch_dimadd_class_tokenclass_token_lennum_patches_per_dim_hnum_patches_per_dim_w num_patchesrr_/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/vlm/llava_next/model/utils.pyget_image_sequence_lengthsrcCsddl}t|ttfs#t|tj|jfstdt|dd| }|\}}|j dd\}}||}||}||kr]||} t || } || d} |dd| || ddf} | S||} t || } || d} |dddd| || f} | S)a\ Unpads a PyTorch tensor of a padded and resized image. Args: tensor (`torch.Tensor`): The image tensor, assumed to be of shape (num_channels, height, width). original_size (`tuple`): The original size of the image (height, width). Returns: `torch.Tensor`: The unpadded image tensor. rNimage_size invalid type: z not valid z2should be either list, tuple, np.ndarray or tensor) numpy isinstancelisttupletorchTensorndarray TypeErrortypetolistshapeint)tensor original_sizenporiginal_heightoriginal_widthcurrent_height current_widthoriginal_aspect_ratiocurrent_aspect_ratio scale_factor new_heightpaddingunpadded_tensor new_widthrrr unpad_images.     r,rpossible_resolutionsreturncCs|\}}d}d}td}|D]=\}}t||||} t|| t|| } } t| | ||} ||| } | |ksC| |krK| |krK| }| }||f}q|S)a Selects the best resolution from a list of possible resolutions based on the original size. This is done by calculating the effective and wasted resolution for each possible resolution. The best fit resolution is the one that maximizes the effective resolution and minimizes the wasted resolution. Args: original_size (tuple): The original size of the image in the format (height, width). possible_resolutions (list): A list of possible resolutions in the format [(height1, width1), (height2, width2), ...]. Returns: tuple: The best fit resolution in the format (height, width). Nrinf)floatminr)rr-r!r"best_fitmax_effective_resolutionmin_wasted_resolutionheightwidthscaledownscaled_widthdownscaled_heighteffective_resolutionwasted_resolutionrrr select_best_resolutionGs    r<cCstddl}t|ts tdt|ttfs+t|tj|jfs'tdt|d| }t ||\}}||||fS)a- Calculate the shape of the image patch grid after the preprocessing for images of any resolution. Args: image_size (`tuple`): The size of the input image in the format (width, height). grid_pinpoints (`List`): A list containing possible resolutions. Each item in the list should be a tuple or list of the form `(height, width)`. patch_size (`int`): The size of each image patch. Returns: tuple: The shape of the image patch grid in the format (width, height). rNz2grid_pinpoints should be a list of tuples or listsrz> not valid, should be either list, tuple, np.ndarray or tensor) rrrrrrrrrrr<) image_sizegrid_pinpoints patch_sizer r5r6rrr get_anyres_image_grid_shapeis r@cCsddlm}|}g}g}t|D]\}} | jddkr| d} | dd} |jj|jj} } |dkr9| | } n |dkrC| | d} | | jdkrNtdt|||j |jj\}}| ||| | d} | d dd dd  } | dd  d d } t| ||} |durtj| |ddddfjg| jdddR| jfdd } | dd dd} tj| | fdd } n| d} |durtj| |d| fdd } || || dqtj|dd }tj|tj|jd }||fS)a Reshape, unpad and then pack each image_feature into a single image_features tensor containing all visual vectors. Args: image_features (`List[torch.Tensor]` of length num_images, each of shape `(num_patches, image_length, embed_dim)`) List of image feature tensor, each contains all the visual feature of all patches. image_sizes (`torch.Tensor` of shape `(num_images, 2)`) Actual image size of each images (H, W). vision_feature_select_strategy (`str`) The feature selection strategy used to select the vision feature from the vision backbone. image_newline (`torch.Tensor` of shape `(embed_dim)`) New line embedding vector. Returns: image_features (`torch.Tensor` of shape `(all_feat_len, embed_dim)`) feature_lens (`List[int]`) token length of each image in image_features r)LlavaNextConfigrNdefaultfullzs. (" %C "