Views
No views yet
patrickamadeus/select-distill-460m-4000-student.models.nanovlm.VisionLanguageModel.
It is not a Transformers AutoModel checkpoint. Use it from a checkout of the
PrefixVLM codebase that defines models.nanovlm.VisionLanguageModel.1git clone https://github.com/patrickamadeus/PrefixVLM.git
2cd PrefixVLM
3pip install -r requirements.txt1from models.nanovlm import VisionLanguageModel
2
3model = VisionLanguageModel.from_pretrained("patrickamadeus/nanovlm-460m-distill_lang")
4model.eval()1python generate_nanovlm.py \
2 --checkpoint patrickamadeus/nanovlm-460m-distill_lang \
3 --image ./assets/cat.png \
4 --prompt "What is in the image?" \
5 --greedy \
6 --max_new_tokens 641import math
2import torch
3from einops import rearrange
4from PIL import Image
5from torchvision.transforms.functional import InterpolationMode, resize, to_tensor
6from transformers import AutoTokenizer
7
8from models.nanovlm import VisionLanguageModel
9
10
11def load_tokenizer(cfg):
12 tokenizer = AutoTokenizer.from_pretrained(
13 cfg.lm_tokenizer,
14 use_fast=True,
15 extra_special_tokens=cfg.vlm_extra_tokens,
16 chat_template=cfg.lm_chat_template,
17 )
18 tokenizer.pad_token = tokenizer.eos_token
19 return tokenizer
20
21
22def resize_to_patch_grid(image, patch_size, max_side_len, resize_to_max_side_len=False):
23 width, height = image.size
24 long_side, short_side = (width, height) if width >= height else (height, width)
25 target_long = max_side_len if resize_to_max_side_len else min(
26 max_side_len,
27 math.ceil(long_side / patch_size) * patch_size,
28 )
29 scale = target_long / long_side
30 target_short = max(patch_size, math.ceil(short_side * scale / patch_size) * patch_size)
31 new_height, new_width = (
32 (target_short, target_long) if width >= height else (target_long, target_short)
33 )
34 return resize(image, [new_height, new_width], interpolation=InterpolationMode.BICUBIC)
35
36
37def split_global_and_tiles(image_tensor, tile_size):
38 if image_tensor.ndim == 3:
39 image_tensor = image_tensor.unsqueeze(0)
40 _, _, height, width = image_tensor.shape
41 if height % tile_size or width % tile_size:
42 raise ValueError(f"image size {(height, width)} is not divisible by {tile_size}")
43
44 n_h, n_w = height // tile_size, width // tile_size
45 tiles = rearrange(
46 image_tensor,
47 "b c (nh ph) (nw pw) -> (b nh nw) c ph pw",
48 ph=tile_size,
49 pw=tile_size,
50 )
51 if (n_h, n_w) == (1, 1):
52 return tiles, (n_h, n_w)
53 global_tile = resize(image_tensor, [tile_size, tile_size])
54 return torch.cat([global_tile, tiles], dim=0), (n_h, n_w)
55
56
57def build_image_string(tokenizer, grid, image_token_length):
58 n_h, n_w = grid
59 text = ""
60 if hasattr(tokenizer, "global_image_token"):
61 text += tokenizer.global_image_token
62 text += tokenizer.image_token * image_token_length
63 if (n_h, n_w) == (1, 1):
64 return text
65
66 for row in range(n_h):
67 for col in range(n_w):
68 text += getattr(tokenizer, f"r{row + 1}c{col + 1}")
69 text += tokenizer.image_token * image_token_length
70 return text
71
72
73def build_inputs(model, image_path, prompt, device):
74 cfg = model.cfg
75 tokenizer = load_tokenizer(cfg)
76 image = Image.open(image_path).convert("RGB")
77 image = resize_to_patch_grid(
78 image,
79 patch_size=cfg.vit_img_size,
80 max_side_len=cfg.max_img_size,
81 resize_to_max_side_len=cfg.resize_to_max_side_len,
82 )
83 image_tensor, grid = split_global_and_tiles(to_tensor(image), cfg.vit_img_size)
84
85 if not hasattr(tokenizer, "global_image_token") and grid[0] * grid[1] == image_tensor.size(0) - 1:
86 image_tensor = image_tensor[1:]
87
88 image_text = build_image_string(tokenizer, grid, cfg.mp_image_token_length)
89 messages = [{"role": "user", "content": image_text + prompt}]
90 prompt_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True)
91 input_ids = torch.tensor(prompt_ids, dtype=torch.long, device=device).unsqueeze(0)
92 attention_mask = torch.ones_like(input_ids)
93 return tokenizer, input_ids, attention_mask, [image_tensor]
94
95
96model_id = "patrickamadeus/nanovlm-460m-distill_lang"
97device = "cuda" if torch.cuda.is_available() else "cpu"
98
99model = VisionLanguageModel.from_pretrained(model_id).to(device)
100model.eval()
101
102tokenizer, input_ids, attention_mask, images = build_inputs(
103 model,
104 "./assets/cat.png",
105 "What is in the image?",
106 torch.device(device),
107)
108
109with torch.inference_mode():
110 output_ids = model.generate(
111 input_ids=input_ids,
112 images=images,
113 attention_mask=attention_mask,
114 max_new_tokens=64,
115 greedy=True,
116 )
117
118print(tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0])