Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,110 @@
+"""Chipoint — image geolocation demo (ZeroGPU).
+
+Full 3-encoder model + density-weighted consensus, on a ZeroGPU (Blackwell) GPU. Weights are pulled from the
+public chiikabu-labs/chipoint repo; DINOv3 comes ungated from ModelScope. Each guess takes a few seconds.
+"""
+import os
+import numpy as np, torch, torch.nn as nn, torch.nn.functional as F
+from pathlib import Path
+from PIL import Image
+import gradio as gr
+import spaces                                              # ZeroGPU
+from huggingface_hub import snapshot_download
+
+DEV = "cuda"                                               # ZeroGPU: load on cuda at module level (emulated until @spaces.GPU)
+M, K = 200, 64
+Gv, LAM, MU, BW1, BW2 = 0.8, 1.7, 0.8, 1492.7, 750.0
+ENC_D = {"so400m256": 1152, "dinov3": 1024, "gopt": 1536}
+
+print("downloading weights ...", flush=True)
+W = Path(snapshot_download("chiikabu-labs/chipoint", allow_patterns=["*.pt", "*.npy"]))
+
+class Head(nn.Module):
+    def __init__(s, D):
+        super().__init__(); s.net = nn.Sequential(nn.Linear(D, 1024), nn.GELU(), nn.Dropout(0.1), nn.Linear(1024, 512))
+    def forward(s, x): return F.normalize(s.net(x), dim=-1)
+
+def load_head(tag, D):
+    h = Head(D); sd = torch.load(W / f"proj_head_{tag}.pt", map_location="cpu")
+    h.load_state_dict({k: v for k, v in sd.items() if k.startswith("net.")})
+    return h.to(DEV).eval()
+
+print("loading heads + galleries + encoders ...", flush=True)
+HEADS = {t: load_head(t, D) for t, D in ENC_D.items()}
+GAL = {t: np.load(W / f"proj_head_gallery_512_{t}.npy", mmap_mode="r") for t in ENC_D}   # mmap, streamed to GPU per query
+GC = np.load(W / "gallery_all_C_so400m256.npy").astype("float32")
+def cell_key(ll, C): return (np.round(ll[..., 0]/C).astype(np.int64) << 20) ^ (np.round(ll[..., 1]/C).astype(np.int64) & 0xFFFFF)
+DENS = {C: dict(zip(*[a.tolist() for a in np.unique(cell_key(GC, C), return_counts=True)])) for C in (0.25, 0.5, 1.0, 2.0)}
+
+import open_clip, torchvision.transforms as T
+from transformers import AutoModel, AutoImageProcessor
+from modelscope import snapshot_download as ms_snapshot
+_OC = {}
+for tag, name in [("so400m256", "ViT-SO400M-16-SigLIP2-256"), ("gopt", "ViT-gopt-16-SigLIP2-256")]:
+    m, _, pre = open_clip.create_model_and_transforms(name, pretrained="webli"); _OC[tag] = (m.to(DEV).eval(), pre)
+_DINO = ms_snapshot("facebook/dinov3-vitl16-pretrain-lvd1689m")
+_DM = AutoModel.from_pretrained(_DINO).to(DEV).eval()
+_DP = AutoImageProcessor.from_pretrained(_DINO)
+_MEAN = torch.tensor(_DP.image_mean).view(1, 3, 1, 1).to(DEV); _STD = torch.tensor(_DP.image_std).view(1, 3, 1, 1).to(DEV)
+_DPRE = T.Compose([T.Resize(256), T.CenterCrop(224), T.ToTensor()])
+
+def zc(a): return (a - a.mean()) / (a.std() + 1e-6)
+
+@torch.no_grad()
+def _encode(tag, pil):
+    if tag == "dinov3":
+        x = ((_DPRE(pil).unsqueeze(0).to(DEV) - _MEAN) / _STD)
+        o = _DM(x); e = o.pooler_output if getattr(o, "pooler_output", None) is not None else o.last_hidden_state[:, 0]
+    else:
+        m, pre = _OC[tag]; e = m.encode_image(pre(pil).unsqueeze(0).to(DEV))
+    return F.normalize(e.float(), dim=-1)
+
+@torch.no_grad()
+def _retrieve(tag, q):                                     # q:[1,512] cuda -> top-M over the gallery
+    g = GAL[tag]; step = 500000
+    bs = torch.full((M,), -1e9, device=DEV); bi = torch.zeros(M, dtype=torch.long, device=DEV)
+    for i in range(0, g.shape[0], step):
+        gg = torch.from_numpy(np.asarray(g[i:i+step], dtype="float16")).to(DEV)
+        s = (q.half() @ gg.T)[0].float(); cs, ci = s.topk(min(M, s.shape[0]))
+        al = torch.cat([bs, cs]); ai = torch.cat([bi, ci + i]); o = al.topk(M).indices
+        bs, bi = al[o], ai[o]; del gg, s
+    return bi.cpu().numpy(), bs.cpu().numpy()
+
+@spaces.GPU(duration=120)
+def locate(pil):
+    pil = pil.convert("RGB"); sc = {}
+    for t in ENC_D:
+        q = F.normalize(HEADS[t](_encode(t, pil)), dim=-1)
+        idx, sim = _retrieve(t, q); z = zc(sim)
+        for j, zz in zip(idx, z): sc[int(j)] = sc.get(int(j), 0.0) + float(zz)
+    ci = np.array(sorted(sc, key=lambda k: -sc[k])[:K], np.int64); csim = np.array([sc[i] for i in ci], "float32")
+    cc = GC[ci]
+    DW = np.mean([[-np.log(DENS[C].get(int(k), 1)+1.0) for k in cell_key(cc, C)] for C in (0.25,0.5,1.0,2.0)], 0)
+    la, lo = np.radians(cc[:, 0]), np.radians(cc[:, 1])
+    D = 2*6371.0*np.arcsin(np.sqrt(np.clip(np.sin((la[:,None]-la[None,:])/2)**2 +
+        np.cos(la)[:,None]*np.cos(la)[None,:]*np.sin((lo[:,None]-lo[None,:])/2)**2, 0, 1)))
+    vote = np.maximum(csim, 0)*np.exp(Gv*zc(DW))
+    KV1 = (vote[None,:]*np.exp(-D/BW1)).sum(1); KV2 = (vote[None,:]*np.exp(-D/BW2)).sum(1)
+    score = zc(csim) + 0.3*zc(DW) + LAM*zc(KV1) + MU*zc(KV2)
+    lat, lon = cc[int(score.argmax())]; return float(lat), float(lon)
+
+def predict(img):
+    if img is None: return "Upload a photo first.", ""
+    lat, lon = locate(img)
+    md = f"### 📍 {lat:.4f}, {lon:.4f}\n[Open in Google Maps](https://maps.google.com/?q={lat:.5f},{lon:.5f})"
+    d = 1.5; bbox = f"{lon-d}%2C{lat-d}%2C{lon+d}%2C{lat+d}"
+    iframe = (f'<iframe width="100%" height="360" frameborder="0" '
+              f'src="https://www.openstreetmap.org/export/embed.html?bbox={bbox}&marker={lat}%2C{lon}"></iframe>')
+    return md, iframe
+
+with gr.Blocks(title="Chipoint") as demo:
+    gr.Markdown("# Chipoint — where was this photo taken?\n"
+                "Upload an outdoor photo; the model guesses its GPS by matching it against 4.9M geotagged images "
+                "and beats the published state of the art on OSV-5M. It's a region/city model, not street-address.")
+    with gr.Row():
+        inp = gr.Image(type="pil", label="Photo")
+        with gr.Column():
+            out_md = gr.Markdown(); out_map = gr.HTML()
+    gr.Button("Locate", variant="primary").click(predict, inp, [out_md, out_map])
+
+demo.queue(max_size=8).launch()