The encoder files are unchanged; only the decoder (and this card) were added.
1import json, numpy as np, torch
2from PIL import Image
3from huggingface_hub import hf_hub_download
4from open_clip.factory import create_vision_encoder_and_transforms
5from modeling_openvision2_decoder import OpenVision2TextDecoder, OpenVision2TextDecoderConfig
6from safetensors.torch import load_file
7
8repo = "UCSC-VLAA/openvision2-vit-huge-patch14-224-vision-only"
9enc = create_vision_encoder_and_transforms(model_name=f"hf-hub:{repo}").eval()
10
11cfg = json.load(open(hf_hub_download(repo, "text_decoder_config.json")))
12dec = OpenVision2TextDecoder(OpenVision2TextDecoderConfig(
13 width=cfg["width"], depth=cfg["depth"], num_heads=cfg["num_heads"], mlp_dim=cfg["mlp_dim"],
14 vocab_size=cfg["vocab_size"], vision_width=cfg["vision_width"]))
15dec.load_state_dict(load_file(hf_hub_download(repo, "caption_decoder.safetensors"))); dec.eval()
16vocab = [l.rstrip("\n") for l in open(hf_hub_download(repo, "bert_base_vocab_bos_eos.txt"))]
17
18def preprocess(path, res=224):
19 im = Image.open(path).convert("RGB"); w, h = im.size; s = res / min(w, h)
20 im = im.resize((round(w*s), round(h*s)), Image.BILINEAR); w, h = im.size
21 l, t = (w-res)//2, (h-res)//2; im = im.crop((l, t, l+res, t+res))
22 x = (np.asarray(im, np.float32) - np.array([.485,.456,.406])*255) / (np.array([.229,.224,.225])*255)
23 return torch.tensor(x.transpose(2,0,1)[None], dtype=torch.float32)
24
25with torch.no_grad():
26 _, tokens = enc(preprocess("image.jpg"))
27 ids = dec.generate(tokens, max_len=64, bos_id=1, eos_id=2)[0].tolist()
28
29words = []
30for i in ids:
31 if i == 2: break # eos
32 if i in (0, 1, 2): continue # pad / bos / eos
33 tk = vocab[i]
34 if tk.startswith("##"): words[-1] = words[-1] + tk[2:] if words else tk[2:]
35 else: words.append(tk)
36print(" ".join(words))