# HiSpatial-3B-RGB

RGB-only variant of [HiSpatial](https://github.com/microsoft/HiSpatial), built on PaliGemma2-3B-mix-448.

## Usage

```python
import torch
from huggingface_hub import hf_hub_download
from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor

BASE = "google/paligemma2-3b-mix-448"
REPO = "lhzzzzzy/HiSpatial-3B-RGB"

model = PaliGemmaForConditionalGeneration.from_pretrained(BASE, torch_dtype=torch.bfloat16)
state_dict = torch.load(
    hf_hub_download(REPO, "weights.pt"), map_location="cpu", weights_only=True
)
model.load_state_dict(state_dict)
model = model.eval().cuda()

processor = PaliGemmaProcessor.from_pretrained(BASE)
```

Inference:

```python
import cv2

image = cv2.cvtColor(cv2.imread("example.jpg"), cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (448, 448))
prompt = "<image>Which object is closer to the camera, the chair or the table?"

inputs = processor(text=prompt, images=image, return_tensors="pt").to(model.device)
with torch.inference_mode():
    output = model.generate(**inputs, max_new_tokens=100, do_sample=False)

print(processor.decode(output[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True))
```