File size: 7,696 Bytes
47bede8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | import base64
import io
import math
from dataclasses import dataclass
from urllib.request import urlopen
import numpy as np
import torch
from PIL import Image, ImageOps
IMAGE_START, IMAGE_PAD, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(5)
COMPRESS_PAD_TO = 4
@dataclass
class ImageInput:
start: int
patches: torch.Tensor
n_vit_h: int
n_vit_w: int
types: torch.Tensor
perm: torch.Tensor
def grid_tokens(best_height, best_width, patch_size, downsample_ratio):
"""Number of LLM tokens the aligner grid occupies (N-layout, incl. row/align padding)."""
n_llm_h = math.ceil((best_height // patch_size) / downsample_ratio)
n_llm_w = math.ceil((best_width // patch_size) / downsample_ratio)
num_tokens = n_llm_h * (n_llm_w + 1) + 2
if n_llm_h % 2 == 1:
num_tokens += n_llm_w + 1
num_tokens += (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2
return n_llm_h, n_llm_w, num_tokens
def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token):
r = height / width
max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5
max_h_float = max_w_float * r
if max_w_float < 1.0:
max_w = 1
max_h = (max_n_token - 2) // (max_w + 1)
if max_h % 2 == 1:
max_h -= 1
best_width = max_w * patch_size * downsample_ratio
best_height = max_h * patch_size * downsample_ratio
elif max_h_float < 2.0:
max_h = 2
max_w = ((max_n_token - 2) // max_h) - 1
assert max_w > 1
best_width = max_w * patch_size * downsample_ratio
best_height = max_h * patch_size * downsample_ratio
else:
max_w = math.floor(max_w_float)
max_h = math.floor(max_h_float)
if max_h % 2 == 1:
max_h -= 1
beta = min(max_w * patch_size * downsample_ratio / width, max_h * patch_size * downsample_ratio / height)
best_width = math.floor(width * beta / patch_size) * patch_size
best_height = math.floor(height * beta / patch_size) * patch_size
n_llm_h, n_llm_w, num_tokens = grid_tokens(best_height, best_width, patch_size, downsample_ratio)
return n_llm_h, n_llm_w, best_height, best_width, num_tokens
def safe_resize(height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token):
max_n_token -= COMPRESS_PAD_TO - 1
n_llm_h, n_llm_w, num_tokens = grid_tokens(best_height, best_width, patch_size, downsample_ratio)
budget = max_n_token
while num_tokens > max_n_token:
n_llm_h, n_llm_w, best_height, best_width, num_tokens = solve_resize_ratio(
height, width, patch_size, downsample_ratio, budget)
budget -= 1
return n_llm_h, n_llm_w, best_height, best_width
def load_image_bytes(record) -> bytes:
"""Load image bytes from raw/base64 data, an Anthropic source, URL, or path."""
data = record.get("data")
if isinstance(data, bytes):
return data
if isinstance(data, str):
return base64.b64decode(data)
source = record.get("source")
if isinstance(source, dict):
if source.get("data") is not None:
return base64.b64decode(source["data"])
if source.get("url"):
return load_image_bytes({"url": source["url"]})
url = record.get("url")
if isinstance(url, str) and url:
if url.startswith("data:"):
header, _, payload = url.partition(",")
if ";base64" not in header:
raise ValueError(f"Unsupported data URL encoding: {header}")
return base64.b64decode(payload)
if url.startswith(("http://", "https://")):
with urlopen(url, timeout=30) as response:
return response.read()
with open(url, "rb") as file:
return file.read()
raise ValueError(f"Cannot load image from record: {list(record.keys())}")
def load_image(record, args):
"""Load and transform one image record into ViT patches."""
p = args.vision_patch_size
with Image.open(io.BytesIO(load_image_bytes(record))) as source:
image = source.convert("RGB")
width, height = image.size
if args.vision_max_wh_ratio is not None and width > height * args.vision_max_wh_ratio:
width = height * args.vision_max_wh_ratio
if 0 < width * height < args.vision_min_pixels:
ratio = (args.vision_min_pixels / (width * height)) ** 0.5
width = int(width * ratio)
height = int(height * ratio)
best_width = math.ceil(width / p) * p
best_height = math.ceil(height / p) * p
n_llm_h, n_llm_w, best_height, best_width = safe_resize(
height, width, best_height, best_width, p, args.vision_downsample_ratio, args.vision_max_n_token)
n_vit_h, n_vit_w = best_height // p, best_width // p
if args.vision_max_wh_ratio is not None and image.width >= args.vision_max_wh_ratio * image.height:
image = image.resize((best_width, best_height))
else:
image = ImageOps.pad(image, (best_width, best_height), color=(127, 127, 127))
x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255
x = ((x - 0.5) / 0.5).to(torch.bfloat16)
patches = x.reshape(3, n_vit_h, p, n_vit_w, p).permute(1, 3, 0, 2, 4).reshape(n_vit_h * n_vit_w, 3, p, p)
return patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w
def build_image_block(n_llm_h: int, n_llm_w: int, start_pos: int):
"""Builds the N-layout token types (final order) and the aligner-row order for IMAGE slots."""
compress_pad = COMPRESS_PAD_TO - 1 - start_pos % COMPRESS_PAD_TO
pad_h = n_llm_h % 2
rows = n_llm_h + pad_h
row_len = n_llm_w + 1
pad_last = rows // 2 * row_len % 2 * 2
types = torch.tensor(([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h + [IMAGE_PAD] * (row_len * pad_h), dtype=torch.int64)
order = torch.arange(rows * row_len).view(rows // 2, 2, row_len).transpose(1, 2).reshape(-1)
image_idx = torch.full((rows * row_len,), -1, dtype=torch.int64)
image_idx.view(rows, row_len)[:n_llm_h, :n_llm_w] = torch.arange(n_llm_h * n_llm_w).view(n_llm_h, n_llm_w)
perm = image_idx[order]
perm = perm[perm >= 0]
types = torch.cat([
torch.full((compress_pad,), IMAGE_PAD, dtype=torch.int64),
torch.tensor([IMAGE_START]),
types[order],
torch.full((pad_last,), IMAGE_PAD, dtype=torch.int64),
torch.tensor([IMAGE_END]),
])
return types, perm
def prepare_vl_inputs(prompt, images, tokenizer, args):
"""Expand image placeholder tokens into sentinel blocks and ImageInput values."""
from encoding_dsv4 import IMAGE_PLACEHOLDER
image_token_id = tokenizer.convert_tokens_to_ids(IMAGE_PLACEHOLDER)
if image_token_id is None or image_token_id == tokenizer.unk_token_id:
raise ValueError(f"Token not found in tokenizer: {IMAGE_PLACEHOLDER}")
prompt_tokens = tokenizer.encode(prompt)
num_placeholders = sum(token == image_token_id for token in prompt_tokens)
if num_placeholders != len(images):
raise ValueError(
f"Found {num_placeholders} image tokens but got {len(images)} images")
tokens, image_inputs = [], []
image_iter = iter(images)
for tok in prompt_tokens:
if tok != image_token_id:
tokens.append(tok)
continue
patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w = load_image(
next(image_iter), args)
types, perm = build_image_block(n_llm_h, n_llm_w, len(tokens))
image_inputs.append(ImageInput(len(tokens), patches, n_vit_h, n_vit_w, types, perm))
tokens += (args.vocab_size + types).tolist()
if not image_inputs:
return tokens, None
return tokens, image_inputs
|