app.py

"""Real pinned regional priors: CPU default plus optional StreetCLIP ZeroGPU.

Model acquisition runs ONLY in the declared HF Space, never on a developer host.
No map-service credentials or paid hardware changes are used. Optional GPU calls use ordinary ZeroGPU quotas.
"""
import os
os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
# Import spaces before Gradio/Torch, as required by HF Spaces runtime integration.
import spaces  # noqa: F401,E402

import importlib.metadata
import copy
from contextlib import nullcontext
from spaces.config import Config as SpacesConfig
import json
from pathlib import Path
import threading
import time
import traceback
import warnings

import gradio as gr
from PIL import Image, ImageOps
import torch
import torch.nn.functional as F

from contracts import (MAX_IMAGE_PIXELS, caption_rows, coordinate_rows, top_k_value)

TARGET_SPACE = "SammySnake/CLIP"
STREET_REPO = "geolocal/StreetCLIP"
STREET_REVISION = "e3561ba2ad9bf14c9efd6b0092607b8497efbfea"
GEO_REPO = "VicenteVivanco/geoclip-geolocation"
GEO_REVISION = "40a931123d2c6451ef7da1d37cf68cded0ec4d35"
BACKBONE_REPO = "openai/clip-vit-large-patch14"
BACKBONE_REVISION = "32bd64288804d66eefd0ccbe215aa642df71cc41"

Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
warnings.filterwarnings("error", category=Image.DecompressionBombWarning)


def read_limit(path):
    try:
        return Path(path).read_text().strip()
    except OSError:
        return None


def cpu_threads():
    available = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)
    quota = read_limit("/sys/fs/cgroup/cpu.max")
    if quota:
        parts = quota.split()
        if len(parts) == 2 and parts[0] != "max":
            available = min(available, max(1, int(parts[0]) // int(parts[1])))
    return max(1, min(4, available))


torch.set_num_threads(cpu_threads())
# The HF/ZeroGPU runtime initializes inter-op scheduling before this app loads.
# Do not reset that process-global one-shot setting after initialization.

_MODELS = {}
_GALLERY = None
_STATE_LOCK = threading.Lock()
_INFER_LOCK = threading.Lock()
_STATE = {
    "streetclip": {"phase": "not_started", "repo": STREET_REPO, "revision": STREET_REVISION},
    "streetclip_gpu": {"phase": "not_started", "repo": STREET_REPO, "revision": STREET_REVISION},
    "geoclip": {"phase": "not_started", "repo": GEO_REPO, "revision": GEO_REVISION},
    "geoclip_gallery": {"phase": "not_started", "rows_encoded": 0},
}


def phase(name, state, **metadata):
    with _STATE_LOCK:
        _STATE[name].update(phase=state, **metadata)
    print(json.dumps({"event": "model_state", "model": name, "phase": state, **metadata}), flush=True)


def features(value):
    """Transformers 4.x tensor / 5.x projected pooler_output compatibility."""
    if isinstance(value, torch.Tensor):
        return value
    if isinstance(getattr(value, "pooler_output", None), torch.Tensor):
        return value.pooler_output
    raise TypeError("Unexpected CLIP feature output; model inference not accepted.")


def load_models():
    """Cloud-only downloader. Failure is exposed, never replaced with dummy weights."""
    global _GALLERY
    if os.environ.get("SPACE_ID") != TARGET_SPACE:
        for name in ("streetclip", "streetclip_gpu", "geoclip", "geoclip_gallery"):
            phase(name, "disabled_outside_target_space")
        return
    start = time.perf_counter()
    try:
        phase("streetclip", "loading")
        from huggingface_hub import snapshot_download
        from transformers import CLIPModel, CLIPProcessor
        folder = snapshot_download(
            STREET_REPO, revision=STREET_REVISION,
            allow_patterns=["*.json", "merges.txt", "pytorch_model.bin"],
        )
        processor = CLIPProcessor.from_pretrained(folder, local_files_only=True)
        model = CLIPModel.from_pretrained(
            folder, local_files_only=True, use_safetensors=False,
        ).eval().to("cpu")
        model.requires_grad_(False)
        _MODELS["streetclip"] = (model, processor)
        phase("streetclip", "ready", load_seconds=round(time.perf_counter() - start, 3), device="cpu")
    except Exception as exc:
        phase("streetclip", "error", error_type=type(exc).__name__)
        traceback.print_exc()

    # Stage an independent copy before decorator registration / Gradio startup.
    # CPU parameters are never moved to CUDA or mutated by a GPU request.
    if SpacesConfig.zero_gpu and "streetclip" in _MODELS:
        try:
            phase("streetclip_gpu", "loading")
            cpu_model, processor = _MODELS["streetclip"]
            gpu_model = copy.deepcopy(cpu_model).to("cuda").eval()
            if next(gpu_model.parameters()).device.type != "cuda":
                raise RuntimeError("StreetCLIP GPU copy was not placed on CUDA.")
            _MODELS["streetclip_gpu"] = (gpu_model, processor)
            phase("streetclip_gpu", "ready", device="cuda", allocation="zerogpu_staged",
                  note="Real GPU execution occurs only during the explicit GPU endpoint call.")
        except Exception as exc:
            phase("streetclip_gpu", "error", error_type=type(exc).__name__)
            traceback.print_exc()
    else:
        phase("streetclip_gpu", "unavailable", reason="ZeroGPU inactive or CPU model failed to load")

    start = time.perf_counter()
    try:
        phase("geoclip", "loading")
        from geoclip import GeoCLIP
        model = GeoCLIP.from_pretrained(GEO_REPO, revision=GEO_REVISION, device="cpu")
        model.requires_grad_(False)
        _MODELS["geoclip"] = model
        phase("geoclip", "ready", load_seconds=round(time.perf_counter() - start, 3), device="cpu",
              gallery_size=len(model.gps_gallery))
        phase("geoclip_gallery", "encoding")
        gallery_features = []
        with torch.inference_mode():
            for first in range(0, len(model.gps_gallery), 1024):
                batch = model.gps_gallery[first:first + 1024]
                gallery_features.append(F.normalize(model.location_encoder(batch), dim=-1))
                with _STATE_LOCK:
                    _STATE["geoclip_gallery"]["rows_encoded"] = min(first + 1024, len(model.gps_gallery))
        _GALLERY = torch.cat(gallery_features, dim=0)
        phase("geoclip_gallery", "ready", rows_encoded=len(_GALLERY), dimension=_GALLERY.shape[1])
    except Exception as exc:
        name = "geoclip_gallery" if "geoclip" in _MODELS else "geoclip"
        phase(name, "error", error_type=type(exc).__name__)
        traceback.print_exc()


# Synchronous root-module setup is intentional for ZeroGPU model staging.
# Local imports / schema tests never enter this branch and never fetch weights.
if os.environ.get("SPACE_ID") == TARGET_SPACE:
    load_models()


def installed_version(name):
    try:
        return importlib.metadata.version(name)
    except importlib.metadata.PackageNotFoundError:
        return "not_installed"


def status():
    with _STATE_LOCK:
        state = json.loads(json.dumps(_STATE))
    return {
        "app": "GeoCLIP + StreetCLIP regional priors", "app_version": 3,
        "space": TARGET_SPACE, "execution": "cpu_default_optional_zerogpu", "models": state,
        "all_models_ready": all(state[name]["phase"] == "ready" for name in
                                ("streetclip", "geoclip", "geoclip_gallery")),
        "optional_gpu_ready": state["streetclip_gpu"]["phase"] == "ready",
        "gpu_runtime_policy": {"endpoint": "/streetclip_score_gpu", "duration_seconds": 30,
                               "uses_normal_zerogpu_quota": True, "automatic_fallback": False},
        "limits": {"upload_bytes": 10_000_000, "image_pixels": MAX_IMAGE_PIXELS,
                   "captions": 100, "coordinate_candidates": 100, "top_k": 20,
                   "queue": 8, "inference_concurrency": 1},
        "runtime": {"torch_threads": torch.get_num_threads(),
                    "torch_interop_threads": torch.get_num_interop_threads(),
                    "cpu_affinity_count": len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else None,
                    "cpu_max": read_limit("/sys/fs/cgroup/cpu.max"),
                    "memory_max": read_limit("/sys/fs/cgroup/memory.max"),
                    "memory_current": read_limit("/sys/fs/cgroup/memory.current")},
        "versions": {name: installed_version(name) for name in
                     ("torch", "transformers", "geoclip", "gradio", "huggingface-hub")},
        "backbone": {"repo": BACKBONE_REPO, "revision": BACKBONE_REVISION},
        "warning": "Scores rank hypotheses, not calibrated location confidence. No precise-location verification.",
    }


def image_rgb(image):
    if not isinstance(image, Image.Image):
        raise gr.Error("Upload a real ground-level image.")
    if image.width < 16 or image.height < 16 or image.width * image.height > MAX_IMAGE_PIXELS:
        raise gr.Error("Image must be at least 16×16 and no larger than 16 megapixels.")
    return ImageOps.exif_transpose(image).convert("RGB")


def validated(parser, value):
    try:
        return parser(value)
    except ValueError as exc:
        raise gr.Error(str(exc)) from exc


def require(name):
    if name not in _MODELS:
        with _STATE_LOCK:
            current = _STATE[name]["phase"]
        raise gr.Error(f"{name} is {current}. Check /status and retry after loading completes.")
    return _MODELS[name]


def common_result(model, image, started):
    return {
        "model": model, "device": "cpu", "elapsed_seconds": round(time.perf_counter() - started, 4),
        "input_width": image.width, "input_height": image.height,
        "coordinate_verified": False, "calibrated_confidence": False,
        "warning": "Regional hypothesis ranking only. Verify with independent geometry/imagery; scores are not probabilities of correctness.",
    }


def _streetclip_score_impl(image, captions_json, top_k, *, model_key, device):
    """Shared real model math; device is controlled by endpoint, not user input."""
    started = time.perf_counter()
    image = image_rgb(image)
    rows = validated(caption_rows, captions_json)
    top_k = validated(top_k_value, top_k)
    model, processor = require(model_key)
    # Do not carry a parent thread lock into ZeroGPU's isolated worker process.
    guard = _INFER_LOCK if device == "cpu" else nullcontext()
    with guard, torch.inference_mode():
        pixels = processor(images=image, return_tensors="pt")["pixel_values"].to(device)
        image_features = F.normalize(features(model.get_image_features(pixel_values=pixels)), dim=-1)
        all_scores = []
        for first in range(0, len(rows), 16):
            texts = [row["text"] for row in rows[first:first + 16]]
            inputs = processor.tokenizer(texts, padding=True, truncation=False, return_tensors="pt")
            if inputs["input_ids"].shape[1] > model.config.text_config.max_position_embeddings:
                raise gr.Error("A caption exceeds the CLIP 77-token limit. Shorten it; no silent truncation is performed.")
            inputs = {key: value.to(device) for key, value in inputs.items()}
            text_features = F.normalize(features(model.get_text_features(**inputs)), dim=-1)
            all_scores.append((image_features @ text_features.T).squeeze(0))
        cosine = torch.cat(all_scores)
        logits = cosine * model.logit_scale.exp()
        # Only CPU scalars cross the ZeroGPU process boundary.
        cosine, logits = cosine.float().cpu(), logits.float().cpu()
        order = torch.argsort(cosine, descending=True)[:min(top_k, len(rows))]
        ranking = [{"rank": rank + 1, **rows[index], "cosine_similarity": float(cosine[index]),
                    "logit": float(logits[index])} for rank, index in enumerate(order.tolist())]
    result = common_result("StreetCLIP", image, started)
    result.update(device=device, revision=STREET_REVISION, candidate_count=len(rows), ranking=ranking,
                  preprocessing="Official 336px CLIP image processor; English text, maximum 77 tokens.",
                  limitations=["StreetCLIP's training dataset excludes China and India.",
                               "Candidate phrasing and candidate coverage affect rankings; cosine is not calibrated confidence."])
    if device == "cuda":
        result["gpu"] = {"name": torch.cuda.get_device_name(),
                         "allocated_bytes": torch.cuda.memory_allocated(),
                         "duration_seconds": 30, "uses_normal_zerogpu_quota": True}
    return result


def streetclip_score(image, captions_json, top_k=10):
    """Default CPU scoring: no GPU allocation is requested by this endpoint."""
    return _streetclip_score_impl(image, captions_json, top_k, model_key="streetclip", device="cpu")


@spaces.GPU(duration=30)
def streetclip_score_gpu(image, captions_json, top_k=10):
    """Optional genuine CUDA inference; ordinary quota/capacity errors propagate."""
    # Match CPU FP32 scoring: no autocast or TF32 approximation.
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False
    return _streetclip_score_impl(image, captions_json, top_k, model_key="streetclip_gpu", device="cuda")


def geoclip_candidates(image, coordinates_json, top_k=10):
    """Score caller-supplied GPS lat/lon hypotheses with the real GeoCLIP encoder."""
    started = time.perf_counter()
    image = image_rgb(image)
    rows = validated(coordinate_rows, coordinates_json)
    top_k = validated(top_k_value, top_k)
    model = require("geoclip")
    with _INFER_LOCK, torch.inference_mode():
        pixels = model.image_encoder.preprocess_image(image)
        coordinates = torch.tensor([[row["lat"], row["lon"]] for row in rows], dtype=torch.float32)
        logits = model(pixels, coordinates).squeeze(0)
        cosine = logits / model.logit_scale.exp()
        order = torch.argsort(logits, descending=True)[:min(top_k, len(rows))]
        ranking = [{"rank": rank + 1, **rows[index], "cosine_similarity": float(cosine[index]),
                    "logit": float(logits[index])} for rank, index in enumerate(order.tolist())]
    result = common_result("GeoCLIP", image, started)
    result.update(revision=GEO_REVISION, candidate_count=len(rows), ranking=ranking,
                  coordinate_order="latitude, longitude in degrees; GPS, not GCJ-02/BD-09",
                  preprocessing="Official pinned OpenAI CLIP ViT-L/14 image processor (224px).",
                  warning="Scores are regional priors, not exact GPS verification; candidate coordinates are input hypotheses.")
    return result


def geoclip_predict(image, top_k=10):
    """Retrieve hypotheses from the author's frozen 100,000-coordinate gallery."""
    started = time.perf_counter()
    image = image_rgb(image)
    top_k = validated(top_k_value, top_k)
    model = require("geoclip")
    if _GALLERY is None:
        raise gr.Error("GeoCLIP coordinate gallery is still encoding or failed; inspect /status.")
    with _INFER_LOCK, torch.inference_mode():
        pixels = model.image_encoder.preprocess_image(image)
        image_features = F.normalize(model.image_encoder(pixels), dim=-1)
        cosine = (image_features @ _GALLERY.T).squeeze(0)
        logits = cosine * model.logit_scale.exp()
        relative_softmax = logits.softmax(dim=-1)
        order = torch.argsort(logits, descending=True)[:top_k]
        ranking = [{"rank": rank + 1, "gallery_index": index,
                    "lat": float(model.gps_gallery[index, 0]), "lon": float(model.gps_gallery[index, 1]),
                    "cosine_similarity": float(cosine[index]), "logit": float(logits[index]),
                    "gallery_relative_softmax": float(relative_softmax[index])}
                   for rank, index in enumerate(order.tolist())]
    result = common_result("GeoCLIP", image, started)
    result.update(revision=GEO_REVISION, gallery_count=len(_GALLERY), ranking=ranking,
                  coordinate_order="latitude, longitude in degrees; GPS, not GCJ-02/BD-09",
                  warning="These are coordinates already present in the author's gallery, not measured camera positions. gallery_relative_softmax is relative to this gallery, not confidence of correctness.")
    return result


def make_demo():
    with gr.Blocks(analytics_enabled=False, delete_cache=(600, 3600)) as demo:
        gr.Markdown("""# GeoCLIP + StreetCLIP · 区域先验检索
这两个模型负责**缩小和排序候选**，不会把照片重建成卫星图，也不证明精确坐标。
默认 CPU 实际推理；模型只下载到此 HF 容器。另提供可选 StreetCLIP GPU 推理，使用正常 ZeroGPU 配额；不承诺无限算力或永久在线。
StreetCLIP 为 **CC-BY-NC-4.0 非商业用途**；其训练数据不包含中国与印度，国内结果必须另行核验。
只上传你有权处理的景物照片。请勿上传用于追踪他人的私人照片。上传文件为临时文件；超过一小时的文件每十分钟清理一次。
""")
        with gr.Tab("StreetCLIP · 图文候选"):
            street_image = gr.Image(type="pil", sources=["upload"], label="原始地面照片（不是 BEV 可视化）")
            captions = gr.Textbox(label='JSON 字符串数组，或 [{"id":"a","text":"..."}]；建议英语', lines=5,
                                  value='["A street in Jiangsu, China.", "A street in Anhui, China.", "A street in Zhejiang, China."]')
            street_k = gr.Number(value=10, precision=0, minimum=1, maximum=20, label="Top K")
            street_button = gr.Button("对候选文字打分")
            street_gpu_button = gr.Button("可选：GPU 打分（消耗正常 ZeroGPU 配额）")
            street_output = gr.JSON(label="相对分数，不是定位置信度")
            street_button.click(streetclip_score, [street_image, captions, street_k], street_output,
                                api_name="streetclip_score", concurrency_id="inference", concurrency_limit=1)
            street_gpu_button.click(streetclip_score_gpu, [street_image, captions, street_k], street_output,
                                    api_name="streetclip_score_gpu", concurrency_id="inference", concurrency_limit=1)
        with gr.Tab("GeoCLIP · 全球粗排"):
            geo_image = gr.Image(type="pil", sources=["upload"], label="原始地面照片")
            geo_k = gr.Number(value=10, precision=0, minimum=1, maximum=20, label="Top K")
            geo_button = gr.Button("检索作者的 100K 坐标候选库")
            geo_output = gr.JSON(label="候选坐标，未经几何核验")
            geo_button.click(geoclip_predict, [geo_image, geo_k], geo_output,
                             api_name="geoclip_predict", concurrency_id="inference", concurrency_limit=1)
        with gr.Tab("GeoCLIP · 自定义区域候选"):
            candidate_image = gr.Image(type="pil", sources=["upload"], label="原始地面照片")
            candidates = gr.Textbox(label='[{"id":"candidate-a","lat":31.0,"lon":120.0}]，最多100项', lines=5,
                                    value='[{"id":"candidate-a","lat":31.0,"lon":120.0}, {"id":"candidate-b","lat":33.0,"lon":118.0}]')
            candidate_k = gr.Number(value=10, precision=0, minimum=1, maximum=20, label="Top K")
            candidate_button = gr.Button("评分给定 GPS 坐标")
            candidate_output = gr.JSON(label="区域先验分数")
            candidate_button.click(geoclip_candidates, [candidate_image, candidates, candidate_k], candidate_output,
                                   api_name="geoclip_candidates", concurrency_id="inference", concurrency_limit=1)
        with gr.Tab("状态 / 来源"):
            status_button = gr.Button("检查真实模型加载状态")
            status_output = gr.JSON(label="状态与固定模型版本")
            status_button.click(status, [], status_output, api_name="status", queue=False)
            gr.Markdown("""[GeoCLIP · NeurIPS 2023](https://github.com/VicenteVivan/geo-clip) ·
[StreetCLIP · model card](https://huggingface.co/geolocal/StreetCLIP) ·
[来源与许可](https://huggingface.co/spaces/SammySnake/CLIP/blob/main/ATTRIBUTION.md)

建议链路：区域先验 → Sample4Geo/其他跨视角模型的候选排序 → OrienterNet BEV 与语义地图匹配 → 独立实景/几何核验。
不要平均这几个模型未经校准的原始分数，也不要把 BEV 特征着色图当作真实卫星图输入它们。
""")
    return demo


if __name__ == "__main__":
    make_demo().queue(max_size=8, default_concurrency_limit=1).launch(
        server_name="0.0.0.0", server_port=7860, max_file_size="10mb", show_error=True,
    )