| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| from __future__ import annotations |
| import torch |
| from .token_constants import DEFAULT_IMAGE_PAD_TOKEN, DEFAULT_IMAGE_TOKEN |
| import numpy as np |
| from PIL import ImageOps |
|
|
| def pad_to_square_and_resize(img, size=1024): |
| padded_img = ImageOps.pad(img, (max(img.size), max(img.size)), color=(0, 0, 0)) |
| resized_img = padded_img.resize((size, size)) |
| return resized_img |
|
|
| def process_images(images, image_processor, model_cfg): |
| result = [] |
| for image in images: |
| pixels = image_processor(images=[image], return_tensors='pt').pixel_values |
| (h, w) = pixels.shape[-2:] |
| result.append((pixels, torch.tensor([[1, h // 16, w // 16]], dtype=torch.long))) |
| return result |
|
|
| def tokenizer_image_token(prompt, tokenizer, image_token_index=-200, return_tensors=None, constant_version=None): |
| prompt = prompt.replace('<image_padding>', '<SPECIAL_23>') |
| ids = torch.tensor(tokenizer(prompt).input_ids, dtype=torch.long) |
| ids[ids == 23] = -500 |
| return ids if return_tensors == 'pt' else ids.tolist() |
|
|
| def compute_num_image_tokens(x, image_aspect_ratio, image_processor): |
| grid = x[0][1] |
| return (int(grid[0, 1]) // 2, int(grid[0, 2]) // 2) |
|
|
| def fix_one_prompt(prompt, config, n_image_tokens): |
| image_idx = 0 |
| while DEFAULT_IMAGE_TOKEN in prompt: |
| newline_token = config.mm_native_newline_token |
| if newline_token: |
| if len(n_image_tokens[image_idx]) == 2: |
| (hh, ww) = n_image_tokens[image_idx] |
| row = DEFAULT_IMAGE_PAD_TOKEN * ww + newline_token |
| rows = row * hh |
| rows = rows.rstrip(newline_token) |
| prompt = prompt.replace('<image>', rows, 1) |
| elif len(n_image_tokens[image_idx]) == 3: |
| (tt, hh, ww) = (int(value) for value in n_image_tokens[image_idx]) |
| row = DEFAULT_IMAGE_PAD_TOKEN * ww + newline_token |
| per_image = row * hh + newline_token + newline_token |
| images = (per_image * tt).rstrip(newline_token + newline_token) |
| prompt = prompt.replace('<image>', images, 1) |
| else: |
| raise NotImplementedError() |
| else: |
| image_tokens = DEFAULT_IMAGE_PAD_TOKEN * np.prod(n_image_tokens[image_idx]) |
| prompt = prompt.replace('<image>', image_tokens, 1) |
| image_idx += 1 |
| return prompt |
|
|
