File size: 2,193 Bytes
e3f7919 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | """SigLIP2-B/16 NaFlex vision tower (frozen), fixed 384x384 -> 24x24 = 576 patches x 768.
encode(images) -> fp16 numpy [N,576,768]: last_hidden_state of the vision tower = AFTER the final post-layernorm,
BEFORE the attention-pooling head (the head is not used). Pixel prep: resize to 384x384 (bicubic, aspect NOT kept),
rescale 1/255, normalise mean=std=0.5, patchify row-major (16x16x3 -> 768, order (ph,pw,c))."""
import os, numpy as np, torch
from PIL import Image
HERE = os.path.dirname(os.path.abspath(__file__))
# Local copy of the encoder if SIGLIP2_PATH is set; otherwise the public Hub repo (auto-downloaded).
MODEL_DIR = os.environ.get("SIGLIP2_PATH") or os.path.join(HERE, "siglip2")
if not os.path.isdir(MODEL_DIR): MODEL_DIR = "google/siglip2-base-patch16-naflex"
S, P = 384, 16; G = S // P # 24
_model = None
def load(device="cpu", dtype=torch.float16):
global _model
from transformers import Siglip2VisionModel
m = Siglip2VisionModel.from_pretrained(MODEL_DIR, dtype=dtype).to(device).eval()
for p in m.parameters(): p.requires_grad_(False)
_model = (m, device, dtype); return _model
def prep(img):
if not isinstance(img, Image.Image): img = Image.open(img)
img = img.convert("RGB").resize((S, S), Image.BICUBIC)
a = (np.asarray(img, dtype=np.float32) / 255.0 - 0.5) / 0.5 # [384,384,3]
a = a.reshape(G, P, G, P, 3).transpose(0, 2, 1, 3, 4).reshape(G * G, P * P * 3) # [576,768]
return a
@torch.no_grad()
def encode(images, device="cpu", batch=8):
"""images: list of PIL images / paths -> np.float16 [N,576,768]"""
global _model
if _model is None or _model[1] != device: load(device)
m, dev, dt = _model; out = []
for i in range(0, len(images), batch):
x = torch.from_numpy(np.stack([prep(im) for im in images[i:i + batch]])).to(dev, dt)
b = x.shape[0]
mask = torch.ones(b, G * G, dtype=torch.long, device=dev)
shapes = torch.tensor([[G, G]] * b, dtype=torch.long, device=dev)
h = m(pixel_values=x, pixel_attention_mask=mask, spatial_shapes=shapes).last_hidden_state
out.append(h.float().cpu().numpy().astype(np.float16))
return np.concatenate(out)
|