Unlike standard VLMs that use many visual tokens per image (e.g., 576 for LLaVA), fVLM compresses each frame to a single visual token using a learned query mechanism:
1import torch
2from torchvision import transforms
3from transformers import AutoTokenizer
4from huggingface_hub import hf_hub_download
5from release.model import FoveatedVLM
6
7# Download checkpoint
8ckpt_path = hf_hub_download("sanps/fVLM-135M", "model.safetensors")
9
10# Build model
11model = FoveatedVLM(
12 llm_name="HuggingFaceTB/SmolLM2-135M-Instruct",
13 dino_name="facebook/dinov2-small",
14 query_dim=384,
15 visual_scale=0.14,
16 deep_query=True,
17)
18
19# Load weights
20state_dict = torch.load(ckpt_path, map_location="cpu")
21model.load_state_dict(state_dict)
22model = model.to("cuda").to(torch.bfloat16).eval()
23
24tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M-Instruct")
25
26# Standard DINO preprocessing
27frame_transform = transforms.Compose([
28 transforms.Resize((224, 224)),
29 transforms.ToTensor(),
30 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
31])
1from PIL import Image
2
3img = Image.open("photo.jpg").convert("RGB")
4frame_tensor = frame_transform(img) # [3, 224, 224]
5frames = frame_tensor.unsqueeze(0).repeat(8, 1, 1, 1) # [8, 3, 224, 224] — replicate to 8
6frames = frames.unsqueeze(0).to("cuda", dtype=torch.bfloat16) # [1, 8, 3, 224, 224]
For video, sample up to 64 frames uniformly. No replication needed.
1# video_frames: list of PIL Images (sampled from video)
2tensors = [frame_transform(f) for f in video_frames]
3frames = torch.stack(tensors).unsqueeze(0).to("cuda", dtype=torch.bfloat16)
4# frames shape: [1, T, 3, 224, 224] where T = number of frames (1-64)
1# Tokenize prompt
2messages = [
3 {"role": "user", "content": "Describe what is happening in this image."},
4]
5text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
6input_ids = tokenizer.encode(text, return_tensors="pt").to("cuda")
7attention_mask = torch.ones_like(input_ids)
8loss_mask = torch.ones_like(input_ids, dtype=torch.float32)
9
10# Forward pass (coarse_fine mode recommended for best quality)
11with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
12 result = model(
13 frames=frames,
14 input_ids=input_ids,
15 attention_mask=attention_mask,
16 loss_mask=loss_mask,
17 mode="coarse_fine",
18 )
19# result["logits"]: [B, S, V] text logits
20# result["loss"]: scalar cross-entropy loss