Fileportz/DeepSeek-V4.1-Flash

1"""Image preprocessing.2 3An image becomes a `n_vit_h x n_vit_w` patch grid for the ViT and a `n_llm_h x n_llm_w` token grid4after the 3x3 aligner downsample, which the LLM sees as5 6    [IMAGE_START] + ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h + [IMAGE_END]7 8Every one of those positions carries `image_token_id` in `input_ids`; only the token type tells them9apart. The IMAGE slots are filled with aligner rows in reading order.10"""11 12import base6413import io14import math15from dataclasses import dataclass16from urllib.request import urlopen17 18import numpy as np19import torch20from PIL import Image, ImageOps21 22TEXT = -123IMAGE_START, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(4)24 25 26@dataclass27class ImageInput:28    start: int29    patches: torch.Tensor30    n_vit_h: int31    n_vit_w: int32    types: torch.Tensor33 34 35def num_image_tokens(n_llm_h: int, n_llm_w: int) -> int:36    return n_llm_h * (n_llm_w + 1) + 237 38 39def llm_grid(best_height: int, best_width: int, patch_size: int, downsample_ratio: int):40    """Token grid the aligner produces from a patch grid of this pixel size."""41    return math.ceil((best_height // patch_size) / downsample_ratio), math.ceil(42        (best_width // patch_size) / downsample_ratio43    )44 45 46def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token):47    """Largest aspect-preserving pixel size whose token grid still fits in max_n_token."""48    r = height / width49    max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.550    max_h_float = max_w_float * r51    cell = patch_size * downsample_ratio52    if max_w_float < 1.0:  # very tall: collapse to a single column53        return (max_n_token - 2) // 2 * cell, cell54    if max_h_float < 1.0:  # very wide: collapse to a single row55        return cell, (max_n_token - 3) * cell56    beta = min(math.floor(max_w_float) * cell / width, math.floor(max_h_float) * cell / height)57    return math.floor(height * beta / patch_size) * patch_size, math.floor(width * beta / patch_size) * patch_size58 59 60def safe_resize(height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token):61    """Shrink the pixel size until the image costs at most max_n_token LLM tokens."""62    n_llm_h, n_llm_w = llm_grid(best_height, best_width, patch_size, downsample_ratio)63    if num_image_tokens(n_llm_h, n_llm_w) > max_n_token:64        best_height, best_width = solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token)65        n_llm_h, n_llm_w = llm_grid(best_height, best_width, patch_size, downsample_ratio)66        assert num_image_tokens(n_llm_h, n_llm_w) <= max_n_token67    return n_llm_h, n_llm_w, best_height, best_width68 69 70def load_image_bytes(record) -> bytes:71    """Load image bytes from raw/base64 data, an Anthropic source, URL, or path."""72    data = record.get("data")73    if isinstance(data, bytes):74        return data75    if isinstance(data, str):76        return base64.b64decode(data)77 78    source = record.get("source")79    if isinstance(source, dict):80        if source.get("data") is not None:81            return base64.b64decode(source["data"])82        if source.get("url"):83            return load_image_bytes({"url": source["url"]})84 85    url = record.get("url")86    if isinstance(url, str) and url:87        if url.startswith("data:"):88            header, _, payload = url.partition(",")89            if ";base64" not in header:90                raise ValueError(f"Unsupported data URL encoding: {header}")91            return base64.b64decode(payload)92        if url.startswith(("http://", "https://")):93            with urlopen(url, timeout=30) as response:94                return response.read()95        with open(url, "rb") as file:96            return file.read()97 98    raise ValueError(f"Cannot load image from record: {list(record.keys())}")99 100 101def plan_image_grid(width: int, height: int, args):102    """Resize plan for an image of the given original size; a pure function of its arguments."""103    p = args.vision_patch_size104    if args.vision_max_wh_ratio is not None and width > height * args.vision_max_wh_ratio:105        width = height * args.vision_max_wh_ratio106    if 0 < width * height < args.vision_min_pixels:107        ratio = (args.vision_min_pixels / (width * height)) ** 0.5108        width = int(width * ratio)109        height = int(height * ratio)110    best_width = math.ceil(width / p) * p111    best_height = math.ceil(height / p) * p112    return safe_resize(height, width, best_height, best_width, p, args.vision_downsample_ratio, args.vision_max_n_token)113 114 115def load_image(record, args):116    """Load and transform one image record into ViT patches."""117    p = args.vision_patch_size118    with Image.open(io.BytesIO(load_image_bytes(record))) as source:119        image = source.convert("RGB")120    n_llm_h, n_llm_w, best_height, best_width = plan_image_grid(image.width, image.height, args)121    n_vit_h, n_vit_w = best_height // p, best_width // p122    if args.vision_max_wh_ratio is not None and image.width >= args.vision_max_wh_ratio * image.height:123        image = image.resize((best_width, best_height))124    else:125        image = ImageOps.pad(image, (best_width, best_height), color=(127, 127, 127))126    x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255127    x = ((x - 0.5) / 0.5).to(torch.bfloat16)128    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)129    return patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w130 131 132def image_token_types(n_llm_h: int, n_llm_w: int) -> torch.Tensor:133    """Default layout: the aligner grid in reading order, one IMAGE_NEW_LINE per row."""134    types = [IMAGE_START]135    types += ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h136    types.append(IMAGE_END)137    return torch.tensor(types, dtype=torch.int64)138 139 140def prepare_vl_inputs(prompt, images, tokenizer, args):141    """Tokenize `prompt`, expanding each image placeholder token into its image span.142 143    Returns (tokens, token_types, image_inputs). Image-span positions carry `args.image_token_id` in144    `tokens` and are distinguished only by `token_types` (TEXT elsewhere). `image_inputs` is None when145    the prompt has no images."""146    from encoding import IMAGE_PLACEHOLDER147 148    # The placeholder is spelled differently across tokenizer revisions, so the id comes from the149    # config; only cross-check it when this tokenizer does know the training-time spelling.150    image_token_id = args.image_token_id151    placeholder_id = tokenizer.convert_tokens_to_ids(IMAGE_PLACEHOLDER)152    if placeholder_id is not None and placeholder_id != tokenizer.unk_token_id:153        assert placeholder_id == image_token_id, (placeholder_id, image_token_id)154    prompt_tokens = tokenizer.encode(prompt)155    num_placeholders = sum(token == image_token_id for token in prompt_tokens)156    if num_placeholders != len(images):157        raise ValueError(f"Found {num_placeholders} image tokens but got {len(images)} images")158    if num_placeholders and not args.vision_enabled:159        raise ValueError("The model config has no vision tower (vision_n_layers == 0) but the prompt contains images")160 161    tokens, token_types, image_inputs = [], [], []162    image_iter = iter(images)163    for tok in prompt_tokens:164        if tok != image_token_id:165            tokens.append(tok)166            token_types.append(TEXT)167            continue168        patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w = load_image(next(image_iter), args)169        types = image_token_types(n_llm_h, n_llm_w)170        image_inputs.append(ImageInput(len(tokens), patches, n_vit_h, n_vit_w, types))171        tokens += [image_token_id] * types.numel()172        token_types += types.tolist()173    return tokens, token_types, image_inputs or None174