import math import torch from typing import Optional from PIL import Image PREDEFINED_RESOLUTIONS = [ (2048, 2048), (2304, 1728), (1728, 2304), (2560, 1440), (1440, 2560), (2496, 1664), (1664, 2496), (3104, 1312), (1312, 3104), (2304, 1792), (1792, 2304), ] def find_closest_resolution(width, height): img_ratio = width / height best_res = None min_diff = float("inf") for w, h in PREDEFINED_RESOLUTIONS: ratio = w / h diff = abs(ratio - img_ratio) if diff < min_diff: min_diff = diff best_res = (w, h) return best_res def resize_pilimage(pil_image, image_size, patch_size=16, resampler=Image.BICUBIC): while min(*pil_image.size) >= 2 * image_size: pil_image = pil_image.resize( tuple(x // 2 for x in pil_image.size), resample=Image.BOX ) m = patch_size width, height = pil_image.width, pil_image.height S_max = image_size * image_size scale = S_max / (width * height) scale = math.sqrt(scale) new_sizes = [ (round(width * scale) // m * m, round(height * scale) // m * m), (round(width * scale) // m * m, math.floor(height * scale) // m * m), (math.floor(width * scale) // m * m, round(height * scale) // m * m), (math.floor(width * scale) // m * m, math.floor(height * scale) // m * m), ] new_sizes = sorted(new_sizes, key=lambda x: x[0] * x[1], reverse=True) for new_size in new_sizes: if new_size[0] * new_size[1] <= S_max: break s1 = width / new_size[0] s2 = height / new_size[1] if s1 < s2: pil_image = pil_image.resize([new_size[0], round(height / s1)], resample=resampler) top = (round(height / s1) - new_size[1]) // 2 pil_image = pil_image.crop((0, top, new_size[0], top + new_size[1])) else: pil_image = pil_image.resize([round(width / s2), new_size[1]], resample=resampler) left = (round(width / s2) - new_size[0]) // 2 pil_image = pil_image.crop((left, 0, left + new_size[0], new_size[1])) return pil_image def calculate_dimensions(max_size, ratio): width = math.sqrt(max_size * max_size * ratio) height = width / ratio width = int(width / 32) * 32 height = int(height / 32) * 32 return width, height def get_rope_index_fix_point( spatial_merge_size, image_token_id, video_token_id, vision_start_token_id, input_ids: Optional[torch.LongTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, skip_vision_start_token=None, fix_point=4096, ) -> tuple[torch.Tensor, torch.Tensor]: if video_grid_thw is not None: video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) video_grid_thw[:, 0] = 1 mrope_position_deltas = [] if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None): total_input_ids = input_ids if attention_mask is None: attention_mask = torch.ones_like(total_input_ids) position_ids = torch.ones( 3, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device, ) image_index, video_index = 0, 0 attention_mask = attention_mask.to(total_input_ids.device) for i, input_ids in enumerate(total_input_ids): input_ids = input_ids[attention_mask[i] == 1] image_nums, video_nums = 0, 0 vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) vision_tokens = input_ids[vision_start_indices + 1] image_nums = (vision_tokens == image_token_id).sum() video_nums = (vision_tokens == video_token_id).sum() input_tokens = input_ids.tolist() llm_pos_ids_list: list = [] st = 0 remain_images, remain_videos = image_nums, video_nums for _ in range(image_nums + video_nums): if image_token_id in input_tokens and remain_images > 0: ed_image = input_tokens.index(image_token_id, st) else: ed_image = len(input_tokens) + 1 if video_token_id in input_tokens and remain_videos > 0: ed_video = input_tokens.index(video_token_id, st) else: ed_video = len(input_tokens) + 1 if ed_image < ed_video: t, h, w = ( image_grid_thw[image_index][0], image_grid_thw[image_index][1], image_grid_thw[image_index][2], ) image_index += 1 remain_images -= 1 ed = ed_image else: t, h, w = ( video_grid_thw[video_index][0], video_grid_thw[video_index][1], video_grid_thw[video_index][2], ) video_index += 1 remain_videos -= 1 ed = ed_video llm_grid_t, llm_grid_h, llm_grid_w = ( t.item(), h.item() // spatial_merge_size, w.item() // spatial_merge_size, ) text_len = ed - st text_len -= skip_vision_start_token[image_index - 1] text_len = max(0, text_len) st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 llm_pos_ids_list.append(torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + st_idx) t_index = torch.arange(llm_grid_t, device=input_ids.device).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() h_index = torch.arange(llm_grid_h, device=input_ids.device).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() w_index = torch.arange(llm_grid_w, device=input_ids.device).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() if skip_vision_start_token[image_index - 1]: if fix_point > 0: fix_point = fix_point - st_idx llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + fix_point + st_idx) fix_point = 0 else: llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) st = ed + llm_grid_t * llm_grid_h * llm_grid_w if st < len(input_tokens): st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 text_len = len(input_tokens) - st llm_pos_ids_list.append(torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + st_idx) llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas else: if attention_mask is not None: position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device) max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0] mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] else: position_ids = ( torch.arange(input_ids.shape[1], device=input_ids.device) .view(1, 1, -1) .expand(3, input_ids.shape[0], -1) ) mrope_position_deltas = torch.zeros( [input_ids.shape[0], 1], device=input_ids.device, dtype=input_ids.dtype, ) return position_ids, mrope_position_deltas