Initial commit

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,81 @@
+---
+title: Sapiens2 CPU
+emoji: 🧍
+colorFrom: indigo
+colorTo: green
+sdk: gradio
+sdk_version: 6.14.0
+app_file: app.py
+pinned: false
+license: other
+---
+
+# Sapiens2 CPU
+
+Meta's `facebook/sapiens2-*` running on free HF CPU. 15 variants exposed: seg, normal, pointmap, pose across 0.4b, 0.8b, 1b, plus seg-5b, normal-5b, pointmap-5b as INT8 ONNX. Curl-callable with a Bearer token.
+
+## Variants and inference time on the included 6000×4000 demo image
+
+| Task | Notes | 0.4b | 0.8b | 1b | 5b (INT8 ONNX) |
+|---|---|---|---|---|---|
+| seg | DOME 29-class body parts | 54 s | 79 s | 142 s | downloads once, then runs |
+| normal | per-pixel surface normals | 50 s | 121 s | 150 s | downloads once, then runs |
+| pointmap | per-pixel XYZ in meters | 76 s | 145 s | 173 s | downloads once, then runs |
+| pose | DETR detect, 308 keypoints | 72 s | 103 s | 174 s | not shipped |
+
+0.4b through 1b run as fp32 PyTorch. 5B runs as INT8 ONNX (5 to 6 GB on disk; fp32 5B would need ~20 GB RAM, more than the free tier provides). Dense tasks share an LRU(2) cache; pose has its own slot and DETR is loaded once. First call per variant downloads then loads (~30-90 s warmup).
+
+The model fixes a 1024×768 input tensor (NCHW with H=1024, W=768, a portrait canvas in Meta's convention). Any input is aspect-preserve resized then padded to that.
+
+## CPU-friendly ONNX exports
+
+Companion repo: [`WeReCooking/sapiens2-onnx`](https://huggingface.co/WeReCooking/sapiens2-onnx) (private). File naming `<task>_<size>_<precision>.onnx` plus a `.onnx.data` external sidecar. 15 ONNX artifacts shipped: 12 covering 0.4b/0.8b/1b (fp16 for seg-0.4b, fp32 for the rest), and 3 new 5B int8 files (seg, normal, pointmap). Cosine similarity vs PyTorch fp32 is 0.999 or better on all shipped variants.
+
+Turnkey CLI built into `app.py` (no sapiens2 / PyTorch dep needed; install `requirements.txt`):
+
+```bash
+export HF_TOKEN=hf_xxx
+python app.py onnx seg 0.4b photo.jpg --output seg_overlay.png
+python app.py onnx normal 1b photo.jpg --output normals.png
+python app.py onnx pointmap 0.8b photo.jpg --output depth.png
+python app.py onnx pose 0.4b photo.jpg --output pose.png
+python app.py onnx seg 5b photo.jpg --output seg_5b.png
+```
+
+## Curl tests
+
+```bash
+TOKEN="hf_xxx"
+SPACE="https://werecooking-sapiens2-cpu.hf.space"
+IMG="https://huggingface.co/spaces/facebook/sapiens2-seg/resolve/main/assets/images/pexels-alex-green-5699868.jpg"
+
+EVT=$(curl -s -X POST "$SPACE/gradio_api/call/predict" \
+  -H "Authorization: Bearer $TOKEN" -H "Content-Type: application/json" \
+  -d "{\"data\":[{\"path\":\"$IMG\",\"meta\":{\"_type\":\"gradio.FileData\"}},\"seg\",\"0.4b\"]}" \
+  | python -c "import sys,json;print(json.load(sys.stdin)['event_id'])")
+curl -sN "$SPACE/gradio_api/call/predict/$EVT" -H "Authorization: Bearer $TOKEN"
+```
+
+## Logs (SSE)
+
+```bash
+curl -N -H "Authorization: Bearer $TOKEN" "https://huggingface.co/api/spaces/WeReCooking/sapiens2-cpu/logs/build"
+curl -N -H "Authorization: Bearer $TOKEN" "https://huggingface.co/api/spaces/WeReCooking/sapiens2-cpu/logs/run"
+```
+
+## 5B INT8 ONNX conversion recipe
+
+The dense 5B variants ship as INT8 ONNX. To re-run the pipeline:
+
+1. Export fp16 ONNX using lazy fp16 init. Call `torch.set_default_dtype(torch.float16)` before `init_model(cfg, None, device="cpu")`, then stream the safetensors file tensor by tensor into the empty fp16 model. This avoids the ~22 GB fp32 init that OOMs on a 32 GB box. Export with `opset_version=18` and no `dynamic_axes`. Force `sys.stdout.reconfigure(encoding="utf-8")` so torch.onnx's success print does not crash on Windows cp1252.
+2. Stream cast fp16 to fp32 on disk via `onnx.external_data_helper.load_external_data_for_model` plus per-tensor `numpy_helper`. Peak RAM stays close to a single tensor (~250 MB). Drop Cast(fp16 / fp32) nodes with transitive rename closure so consumers point at the original input.
+3. Run `quantize.shape_inference.quant_pre_process(skip_onnx_shape=True, skip_optimization=True)`. This routes through ORT symbolic shape inference which understands sapiens2 windowed attention. Vanilla `onnx.shape_inference` errors with `(6144) vs (512)` on the pointmap and normal heads.
+4. `quantize_dynamic(weight_type=QuantType.QInt8, per_channel=True, op_types_to_quantize=["MatMul"], use_external_data_format=True)`. This lowers to `MatMulIntegerToFloat`, which accepts fp32 input and has no 2D-only filter (unlike `MatMulNBitsQuantizer` which silently skips 3D packed-QKV weights).
+
+Pose-5b is not shipped. It uses a different forward signature (single person bbox cropped tensor) and the int8 quantize attempt did not complete on the available hardware.
+
+## Files
+
+* `app.py` everything: Gradio Space UI, PyTorch dispatch for 0.4b/0.8b/1b, ORT for 5B, inlined keypoint visualization, plus the `python app.py onnx ...` CLI
+* `requirements.txt` Python deps including `sapiens @ git+https://github.com/facebookresearch/sapiens2.git`
+* `packages.txt` apt deps (`libgl1`, `libglib2.0-0`) installed by the Gradio SDK at build time

app.py

@@ -0,0 +1,706 @@
+"""Sapiens2 multi-task CPU: seg / normal / pointmap / pose at 0.4b/0.8b/1b plus 5B INT8 ONNX.
+
+5B (seg, normal, pointmap) runs via INT8 ONNX from WeReCooking/sapiens2-onnx; pose-5b not shipped.
+Pose top-down: DETR finds people, sapiens2 estimates 308 keypoints per crop.
+Lazy-load with LRU cache (keeps 2 dense models + 1 pose model resident).
+Per-task API endpoint via Gradio's auto-API (curl-able with Bearer token).
+
+Also exposes a standalone ONNX CLI mode that does not need PyTorch or sapiens2:
+    python app.py onnx seg 0.4b photo.jpg --output seg.png
+    python app.py onnx pointmap 5b photo.jpg --output depth.png
+"""
+# Block mmpretrain: mmdet's reid modules try to import it via try/except ImportError,
+# but mmpretrain raises TypeError on import (transformers API drift) which escapes
+# the except and kills the process.
+import sys
+sys.modules["mmpretrain"] = None
+
+
+# --- ONNX CLI (standalone, no PyTorch/sapiens2 import) ----------------------
+def _onnx_cli():
+    """Run a published sapiens2 ONNX model on a local image. Only needs numpy,
+    onnxruntime, huggingface_hub, opencv-python-headless."""
+    import argparse
+    import os
+    import time
+    from pathlib import Path
+    import numpy as np
+    import cv2
+    import onnxruntime as ort
+    from huggingface_hub import hf_hub_download
+
+    DEFAULT_REPO = "WeReCooking/sapiens2-onnx"
+    PRECISIONS = {("seg", "0.4b"): "fp16"}  # only seg-0.4b is fp16; rest fp32 or int8 for 5B
+    INPUT_HW = (1024, 768)
+
+    parser = argparse.ArgumentParser(prog="app.py onnx")
+    parser.add_argument("task", choices=["seg", "normal", "pointmap", "pose"])
+    parser.add_argument("size", choices=["0.4b", "0.8b", "1b", "5b"])
+    parser.add_argument("image", help="Local image path")
+    parser.add_argument("--cache-dir", default="./onnx_cache")
+    parser.add_argument("--token", default=os.environ.get("HF_TOKEN"))
+    parser.add_argument("--output", default=None, help="Save the visualization here")
+    parser.add_argument("--repo", default=DEFAULT_REPO)
+    args = parser.parse_args(sys.argv[2:])
+
+    precision = PRECISIONS.get((args.task, args.size), "int8" if args.size == "5b" else "fp32")
+    filename = f"{args.task}/{args.task}_{args.size}_{precision}.onnx"
+    print(f"[1/3] downloading {filename} from {args.repo}", flush=True)
+    t0 = time.time()
+    onnx_path = hf_hub_download(repo_id=args.repo, filename=filename, local_dir=args.cache_dir, token=args.token)
+    hf_hub_download(repo_id=args.repo, filename=f"{filename}.data", local_dir=args.cache_dir, token=args.token)
+    print(f"  ready in {time.time()-t0:.1f}s", flush=True)
+
+    img = cv2.imread(args.image, cv2.IMREAD_COLOR)
+    if img is None:
+        raise FileNotFoundError(args.image)
+    H, W = INPUT_HW
+    h0, w0 = img.shape[:2]
+    scale = min(W / w0, H / h0)
+    new_w, new_h = int(round(w0 * scale)), int(round(h0 * scale))
+    resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
+    canvas = np.zeros((H, W, 3), dtype=np.uint8)
+    top = (H - new_h) // 2
+    left = (W - new_w) // 2
+    canvas[top:top + new_h, left:left + new_w] = resized
+    mean = (123.675, 116.28, 103.53)
+    std = (58.395, 57.12, 57.375)
+    x = canvas.astype(np.float32)
+    for c in range(3):
+        x[:, :, c] = (x[:, :, c] - mean[c]) / std[c]
+    x = x.transpose(2, 0, 1)[None]
+
+    print(f"[2/3] ORT forward (input {x.shape} {x.dtype})", flush=True)
+    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+    t0 = time.time()
+    out = sess.run(None, {sess.get_inputs()[0].name: x})
+    print(f"  forward {time.time()-t0:.1f}s, outputs={[o.shape for o in out]}", flush=True)
+
+    print(f"[3/3] postprocess + preview", flush=True)
+    if args.task == "pose":
+        heatmaps = out[0][0]
+        K, hH, hW = heatmaps.shape
+        flat = heatmaps.reshape(K, -1)
+        peak = flat.argmax(axis=1)
+        ys, xs = np.unravel_index(peak, (hH, hW))
+        scores = flat.max(axis=1)
+        inp_y = ys * (INPUT_HW[0] / hH)
+        inp_x = xs * (INPUT_HW[1] / hW)
+        scale_y = h0 / new_h
+        scale_x = w0 / new_w
+        img_y = (inp_y - top) * scale_y
+        img_x = (inp_x - left) * scale_x
+        n_visible = int((scores > 0.3).sum())
+        print(f"  {n_visible}/{K} keypoints above 0.3 confidence (range {scores.min():.3f} to {scores.max():.3f})")
+        if args.output:
+            for i in range(K):
+                if scores[i] < 0.3:
+                    continue
+                cv2.circle(img, (int(img_x[i]), int(img_y[i])), 4, (0, 255, 0), -1)
+            cv2.imwrite(args.output, img)
+            print(f"  saved {args.output}")
+        return
+
+    if args.task == "seg":
+        logits = out[0][0]
+        class_map = logits.argmax(axis=0).astype(np.int32)
+        class_map_crop = class_map[top:top + new_h, left:left + new_w]
+        class_map_full = cv2.resize(class_map_crop, (w0, h0), interpolation=cv2.INTER_NEAREST)
+        classes = np.unique(class_map_full).tolist()
+        print(f"  classes detected: {classes[:15]}")
+        if args.output:
+            palette = (np.random.RandomState(42).rand(29, 3) * 255).astype(np.uint8)
+            cv2.imwrite(args.output, palette[class_map_full])
+            print(f"  saved {args.output}")
+        return
+
+    if args.task == "normal":
+        normal_raw = out[0][0].transpose(1, 2, 0)
+        norm = np.linalg.norm(normal_raw, axis=2, keepdims=True)
+        normal_unit = normal_raw / np.maximum(norm, 1e-8)
+        normal_crop = normal_unit[top:top + new_h, left:left + new_w]
+        normal_full = cv2.resize(normal_crop, (w0, h0), interpolation=cv2.INTER_LINEAR)
+        if args.output:
+            rgb = (((normal_full + 1.0) / 2.0) * 255).clip(0, 255).astype(np.uint8)
+            cv2.imwrite(args.output, rgb)
+            print(f"  saved {args.output}")
+        return
+
+    # pointmap
+    pointmap_rel = out[0][0].transpose(1, 2, 0)
+    s = out[1][0, 0] if len(out) > 1 else 1.0
+    pointmap_metric = pointmap_rel / max(float(s), 1e-8)
+    z = pointmap_metric[..., 2]
+    z_crop = z[top:top + new_h, left:left + new_w]
+    z_full = cv2.resize(z_crop, (w0, h0), interpolation=cv2.INTER_LINEAR)
+    zmin, zmax = float(z_full.min()), float(z_full.max())
+    print(f"  Z range: [{zmin:.2f}, {zmax:.2f}] meters")
+    if args.output:
+        z_norm = ((z_full - zmin) / max(zmax - zmin, 1e-8) * 255).astype(np.uint8)
+        cv2.imwrite(args.output, z_norm)
+        print(f"  saved {args.output}")
+
+
+if len(sys.argv) > 1 and sys.argv[1] == "onnx":
+    _onnx_cli()
+    sys.exit(0)
+
+
+# --- Gradio path -----------------------------------------------------------
+import glob
+import os
+import time
+import traceback
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+from PIL import Image
+
+# --- Catalog ----------------------------------------------------------------
+VARIANTS = {
+    ("seg", "0.4b"):       {"repo": "facebook/sapiens2-seg-0.4b",       "filename": "sapiens2_0.4b_seg.safetensors",       "config_glob": "**/sapiens2_0.4b_seg*shutterstock*1024x768*.py",       "kind": "seg"},
+    ("seg", "0.8b"):       {"repo": "facebook/sapiens2-seg-0.8b",       "filename": "sapiens2_0.8b_seg.safetensors",       "config_glob": "**/sapiens2_0.8b_seg*shutterstock*1024x768*.py",       "kind": "seg"},
+    ("seg", "1b"):         {"repo": "facebook/sapiens2-seg-1b",         "filename": "sapiens2_1b_seg.safetensors",         "config_glob": "**/sapiens2_1b_seg*shutterstock*1024x768*.py",         "kind": "seg"},
+    ("normal", "0.4b"):    {"repo": "facebook/sapiens2-normal-0.4b",    "filename": "sapiens2_0.4b_normal.safetensors",    "config_glob": "**/sapiens2_0.4b_normal*metasim*1024x768*.py",        "kind": "normal"},
+    ("normal", "0.8b"):    {"repo": "facebook/sapiens2-normal-0.8b",    "filename": "sapiens2_0.8b_normal.safetensors",    "config_glob": "**/sapiens2_0.8b_normal*metasim*1024x768*.py",        "kind": "normal"},
+    ("normal", "1b"):      {"repo": "facebook/sapiens2-normal-1b",      "filename": "sapiens2_1b_normal.safetensors",      "config_glob": "**/sapiens2_1b_normal*metasim*1024x768*.py",          "kind": "normal"},
+    ("pointmap", "0.4b"):  {"repo": "facebook/sapiens2-pointmap-0.4b",  "filename": "sapiens2_0.4b_pointmap.safetensors",  "config_glob": "**/sapiens2_0.4b_pointmap*render_people*1024x768*.py", "kind": "pointmap"},
+    ("pointmap", "0.8b"):  {"repo": "facebook/sapiens2-pointmap-0.8b",  "filename": "sapiens2_0.8b_pointmap.safetensors",  "config_glob": "**/sapiens2_0.8b_pointmap*render_people*1024x768*.py", "kind": "pointmap"},
+    ("pointmap", "1b"):    {"repo": "facebook/sapiens2-pointmap-1b",    "filename": "sapiens2_1b_pointmap.safetensors",    "config_glob": "**/sapiens2_1b_pointmap*render_people*1024x768*.py",   "kind": "pointmap"},
+    ("pose", "0.4b"):      {"repo": "facebook/sapiens2-pose-0.4b",      "filename": "sapiens2_0.4b_pose.safetensors",      "config_glob": "**/sapiens2_0.4b_keypoints308*shutterstock_goliath*1024x768*.py", "kind": "pose"},
+    ("pose", "0.8b"):      {"repo": "facebook/sapiens2-pose-0.8b",      "filename": "sapiens2_0.8b_pose.safetensors",      "config_glob": "**/sapiens2_0.8b_keypoints308*shutterstock_goliath*1024x768*.py", "kind": "pose"},
+    ("pose", "1b"):        {"repo": "facebook/sapiens2-pose-1b",        "filename": "sapiens2_1b_pose.safetensors",        "config_glob": "**/sapiens2_1b_keypoints308*shutterstock_goliath*1024x768*.py",   "kind": "pose"},
+    # 5B variants run via prebuilt INT8 ONNX from WeReCooking/sapiens2-onnx.
+    # fp32 5B PyTorch (~20 GB) won't fit in the free CPU Space's 16 GB; INT8 ONNX is ~5-6 GB.
+    # pose-5b is intentionally absent — INT8 wasn't successfully built for it.
+    ("seg", "5b"):         {"onnx_repo": "WeReCooking/sapiens2-onnx", "onnx_filename": "seg/seg_5b_int8.onnx",           "kind": "seg"},
+    ("normal", "5b"):      {"onnx_repo": "WeReCooking/sapiens2-onnx", "onnx_filename": "normal/normal_5b_int8.onnx",     "kind": "normal"},
+    ("pointmap", "5b"):    {"onnx_repo": "WeReCooking/sapiens2-onnx", "onnx_filename": "pointmap/pointmap_5b_int8.onnx", "kind": "pointmap"},
+}
+
+DENSE_KINDS = {"seg", "normal", "pointmap"}
+
+_MODELS: dict = {}       # (task, size) -> dense model (LRU)
+_POSE_MODELS: dict = {}  # (task, size) -> pose model (separate cache so DETR survives)
+_DETECTOR = None         # tuple(processor, model) — lazily loaded once
+_POSE_METAINFO = None
+_ORT_SESSIONS: dict = {} # (task, "5b") -> onnxruntime InferenceSession
+_MAX_CACHED = 2
+_DOME_CLASSES_29 = None
+
+
+_SAPIENS_PKG_ROOT = None
+
+
+def _sapiens_root() -> Path:
+    """Return the directory containing the installed sapiens package."""
+    global _SAPIENS_PKG_ROOT
+    if _SAPIENS_PKG_ROOT is None:
+        import sapiens  # imported lazily because it has side effects (mmdet etc.)
+        _SAPIENS_PKG_ROOT = Path(sapiens.__file__).resolve().parent
+    return _SAPIENS_PKG_ROOT
+
+
+def _find_config(pattern: str) -> str:
+    # cfg_glob comes in as "**/sapiens2_..._1024x768*.py"; rglob applies the leading ** implicitly
+    leaf = pattern.split("/")[-1]
+    root = _sapiens_root()
+    matches = list(root.rglob(leaf))
+    if not matches:
+        raise FileNotFoundError(f"No config matching {leaf} under {root}")
+    return str(matches[0])
+
+
+def _get_dense_model(task: str, size: str):
+    """Lazy-load + LRU-cache for seg/normal/pointmap."""
+    key = (task, size)
+    if key in _MODELS:
+        _MODELS[key] = _MODELS.pop(key)
+        return _MODELS[key]
+
+    spec = VARIANTS[key]
+    from sapiens.dense.models import init_model
+    if spec["kind"] == "normal":
+        from sapiens.dense.models import NormalEstimator  # noqa: F401
+    elif spec["kind"] == "pointmap":
+        from sapiens.dense.models import PointmapEstimator  # noqa: F401
+
+    config = _find_config(spec["config_glob"])
+
+    from huggingface_hub import hf_hub_download
+    local_dir = f"/tmp/sapiens_models/{task}-{size}"
+    os.makedirs(local_dir, exist_ok=True)
+    ckpt = hf_hub_download(repo_id=spec["repo"], filename=spec["filename"], local_dir=local_dir)
+
+    model = init_model(config, ckpt, device="cpu")
+
+    while len(_MODELS) >= _MAX_CACHED:
+        oldest = next(iter(_MODELS))
+        del _MODELS[oldest]
+        import gc
+        gc.collect()
+    _MODELS[key] = model
+    return model
+
+
+def _get_pose_metainfo():
+    global _POSE_METAINFO
+    if _POSE_METAINFO is None:
+        from sapiens.pose.datasets import parse_pose_metainfo
+        meta_cfg = _find_config("**/pose/configs/**/keypoints308.py")
+        import importlib.util
+        spec_obj = importlib.util.spec_from_file_location("keypoints308_meta", meta_cfg)
+        mod = importlib.util.module_from_spec(spec_obj)
+        spec_obj.loader.exec_module(mod)
+        ds_info = getattr(mod, "dataset_info", None)
+        if ds_info is None:
+            raise RuntimeError(f"No dataset_info in {meta_cfg}")
+        _POSE_METAINFO = parse_pose_metainfo(ds_info)
+    return _POSE_METAINFO
+
+
+def _get_pose_model(size: str):
+    key = ("pose", size)
+    if key in _POSE_MODELS:
+        return _POSE_MODELS[key]
+    spec = VARIANTS[key]
+    from sapiens.pose.models import init_model
+    from sapiens.pose.datasets import UDPHeatmap
+
+    config = _find_config(spec["config_glob"])
+    from huggingface_hub import hf_hub_download
+    local_dir = f"/tmp/sapiens_models/pose-{size}"
+    os.makedirs(local_dir, exist_ok=True)
+    ckpt = hf_hub_download(repo_id=spec["repo"], filename=spec["filename"], local_dir=local_dir)
+    model = init_model(config, ckpt, device="cpu")
+
+    codec_cfg = dict(model.cfg.codec)
+    assert codec_cfg.pop("type") == "UDPHeatmap"
+    model.codec = UDPHeatmap(**codec_cfg)
+    model.pose_metainfo = _get_pose_metainfo()
+
+    # Free the largest cached dense model first if more than one pose model present.
+    while len(_POSE_MODELS) >= 1:
+        oldest = next(iter(_POSE_MODELS))
+        del _POSE_MODELS[oldest]
+        import gc
+        gc.collect()
+    _POSE_MODELS[key] = model
+    return model
+
+
+def _get_detector():
+    global _DETECTOR
+    if _DETECTOR is None:
+        import torch  # noqa: F401
+        from transformers import DetrImageProcessor, DetrForObjectDetection
+        proc = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+        det = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50").eval()
+        _DETECTOR = (proc, det)
+    return _DETECTOR
+
+
+def _load_dome_classes():
+    global _DOME_CLASSES_29
+    if _DOME_CLASSES_29 is None:
+        from sapiens.dense.src.datasets.seg.seg_utils import DOME_CLASSES_29
+        _DOME_CLASSES_29 = DOME_CLASSES_29
+    return _DOME_CLASSES_29
+
+
+def _get_padding(data_samples):
+    ds = data_samples[0] if isinstance(data_samples, list) and data_samples else data_samples
+    if hasattr(ds, "padding_size"):
+        return tuple(ds.padding_size)
+    if hasattr(ds, "metainfo") and isinstance(ds.metainfo, dict):
+        if "padding_size" in ds.metainfo:
+            return tuple(ds.metainfo["padding_size"])
+        if "pad_shape" in ds.metainfo and "img_shape" in ds.metainfo:
+            ph, pw = ds.metainfo["pad_shape"][:2]
+            ih, iw = ds.metainfo["img_shape"][:2]
+            return (0, pw - iw, 0, ph - ih)
+    if isinstance(ds, dict):
+        meta = ds.get("meta") or ds
+        if "padding_size" in meta:
+            return tuple(meta["padding_size"])
+    return (0, 0, 0, 0)
+
+
+# --- Per-task inference -----------------------------------------------------
+def _infer_seg(image_bgr, model):
+    import torch
+    import torch.nn.functional as F
+    import cv2
+    h0, w0 = image_bgr.shape[:2]
+    data = model.pipeline(dict(img=image_bgr))
+    data = model.data_preprocessor(data)
+    with torch.no_grad():
+        logits = model(data["inputs"])
+    logits = F.interpolate(logits, size=(h0, w0), mode="bilinear", align_corners=False)
+    label_map = logits.argmax(dim=1).squeeze(0).cpu().numpy().astype(np.int32)
+
+    classes = _load_dome_classes()
+    palette = np.zeros((256, 3), dtype=np.uint8)
+    for cid, meta in classes.items():
+        palette[cid] = meta["color"][::-1]
+    color_mask = palette[label_map]
+    overlay_bgr = cv2.addWeighted(image_bgr, 0.5, color_mask, 0.5, 0)
+    overlay_rgb = cv2.cvtColor(overlay_bgr, cv2.COLOR_BGR2RGB)
+    uniq = sorted(int(c) for c in np.unique(label_map))
+    labels = [classes[c]["name"].replace("_", " ") for c in uniq if c in classes]
+    return Image.fromarray(overlay_rgb), f"classes: {', '.join(labels)}"
+
+
+def _infer_normal(image_bgr, model):
+    import torch
+    data = model.pipeline(dict(img=image_bgr))
+    data = model.data_preprocessor(data)
+    inputs, data_samples = data["inputs"], data["data_samples"]
+    if inputs.ndim == 3:
+        inputs = inputs.unsqueeze(0)
+    with torch.no_grad():
+        normal = model(inputs)
+        normal = normal / normal.norm(dim=1, keepdim=True).clamp_min(1e-8)
+    pl, pr, pt, pb = _get_padding(data_samples)
+    normal = normal[:, :, pt:inputs.shape[2] - pb, pl:inputs.shape[3] - pr]
+    normal_hwc = normal.squeeze(0).cpu().float().numpy().transpose(1, 2, 0)
+    rgb = (((normal_hwc + 1.0) / 2.0) * 255.0).clip(0, 255).astype(np.uint8)
+    return Image.fromarray(rgb), f"normal map {rgb.shape}"
+
+
+def _infer_pointmap(image_bgr, model):
+    import torch
+    data = model.pipeline(dict(img=image_bgr))
+    data = model.data_preprocessor(data)
+    inputs, data_samples = data["inputs"], data["data_samples"]
+    if inputs.ndim == 3:
+        inputs = inputs.unsqueeze(0)
+    with torch.no_grad():
+        out = model(inputs)
+    if isinstance(out, tuple) and len(out) == 2:
+        pointmap, scale = out
+        pointmap = pointmap / scale.clamp_min(1e-8)
+    else:
+        pointmap = out
+    pl, pr, pt, pb = _get_padding(data_samples)
+    pointmap = pointmap[:, :, pt:inputs.shape[2] - pb, pl:inputs.shape[3] - pr]
+    pmap_hwc = pointmap.squeeze(0).cpu().float().numpy().transpose(1, 2, 0)
+    z = pmap_hwc[..., 2]
+    z_min, z_max = float(z.min()), float(z.max())
+    z_norm = (z - z_min) / max(z_max - z_min, 1e-8)
+    z_rgb = (z_norm * 255).astype(np.uint8)
+    rgb = np.stack([z_rgb, z_rgb, z_rgb], axis=-1)
+    return Image.fromarray(rgb), f"pointmap {pmap_hwc.shape} | Z [{z_min:.2f}, {z_max:.2f}]"
+
+
+# --- 5B INT8 ONNX path -------------------------------------------------------
+def _get_ort_session(task: str):
+    """Lazy-load + cache an ORT session for {task}_5b_int8.onnx."""
+    key = (task, "5b")
+    sess = _ORT_SESSIONS.get(key)
+    if sess is not None:
+        return sess
+    import onnxruntime as ort
+    from huggingface_hub import hf_hub_download
+    spec = VARIANTS[key]
+    cache_dir = os.environ.get("ONNX_5B_CACHE", "/app/onnx_5b")
+    os.makedirs(cache_dir, exist_ok=True)
+    # Download both the graph .onnx and its external-data sidecar.
+    fn = spec["onnx_filename"]
+    onnx_path = hf_hub_download(repo_id=spec["onnx_repo"], filename=fn, local_dir=cache_dir)
+    hf_hub_download(repo_id=spec["onnx_repo"], filename=fn + ".data", local_dir=cache_dir)
+    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+    _ORT_SESSIONS[key] = sess
+    return sess
+
+
+def _infer_dense_5b(image_bgr, task: str):
+    """5B inference: preprocess via the 0.4b PyTorch pipeline (cached), forward via ORT INT8."""
+    import torch
+    import torch.nn.functional as F
+    import cv2
+
+    # Use the 0.4b model's pipeline+preprocessor for image prep — it's already in cache for warm calls.
+    proxy = _get_dense_model(task, "0.4b")
+    data = proxy.pipeline(dict(img=image_bgr))
+    data = proxy.data_preprocessor(data)
+    inputs, data_samples = data["inputs"], data["data_samples"]
+    if inputs.ndim == 3:
+        inputs = inputs.unsqueeze(0)
+
+    sess = _get_ort_session(task)
+    out = sess.run(None, {sess.get_inputs()[0].name: inputs.float().cpu().numpy()})
+
+    if task == "seg":
+        logits = torch.from_numpy(out[0])
+        h0, w0 = image_bgr.shape[:2]
+        logits = F.interpolate(logits, size=(h0, w0), mode="bilinear", align_corners=False)
+        label_map = logits.argmax(dim=1).squeeze(0).numpy().astype(np.int32)
+        classes = _load_dome_classes()
+        palette = np.zeros((256, 3), dtype=np.uint8)
+        for cid, meta in classes.items():
+            palette[cid] = meta["color"][::-1]
+        color_mask = palette[label_map]
+        overlay_bgr = cv2.addWeighted(image_bgr, 0.5, color_mask, 0.5, 0)
+        overlay_rgb = cv2.cvtColor(overlay_bgr, cv2.COLOR_BGR2RGB)
+        uniq = sorted(int(c) for c in np.unique(label_map))
+        labels = [classes[c]["name"].replace("_", " ") for c in uniq if c in classes]
+        return Image.fromarray(overlay_rgb), f"classes: {', '.join(labels)}"
+
+    if task == "normal":
+        normal = torch.from_numpy(out[0])
+        normal = normal / normal.norm(dim=1, keepdim=True).clamp_min(1e-8)
+        pl, pr, pt, pb = _get_padding(data_samples)
+        normal = normal[:, :, pt:inputs.shape[2] - pb, pl:inputs.shape[3] - pr]
+        normal_hwc = normal.squeeze(0).numpy().transpose(1, 2, 0)
+        rgb = (((normal_hwc + 1.0) / 2.0) * 255.0).clip(0, 255).astype(np.uint8)
+        return Image.fromarray(rgb), f"normal map {rgb.shape}"
+
+    # pointmap — ONNX produces (pointmap [1,3,H,W], scale [1,1]); divide to recover metric depths.
+    pointmap = torch.from_numpy(out[0])
+    if len(out) > 1:
+        scale = torch.from_numpy(out[1])
+        pointmap = pointmap / scale.clamp_min(1e-8)
+    pl, pr, pt, pb = _get_padding(data_samples)
+    pointmap = pointmap[:, :, pt:inputs.shape[2] - pb, pl:inputs.shape[3] - pr]
+    pmap_hwc = pointmap.squeeze(0).numpy().transpose(1, 2, 0)
+    z = pmap_hwc[..., 2]
+    z_min, z_max = float(z.min()), float(z.max())
+    z_norm = (z - z_min) / max(z_max - z_min, 1e-8)
+    z_rgb = (z_norm * 255).astype(np.uint8)
+    rgb = np.stack([z_rgb, z_rgb, z_rgb], axis=-1)
+    return Image.fromarray(rgb), f"pointmap {pmap_hwc.shape} | Z [{z_min:.2f}, {z_max:.2f}]"
+
+
+# Inlined from the upstream Meta sample (pose keypoint render).
+# Draws skeleton links + colored keypoints; thickness/radius are picked by the caller.
+def visualize_keypoints(
+    image: np.ndarray,
+    keypoints,
+    keypoints_visible,
+    keypoint_scores,
+    *,
+    radius: int = 4,
+    thickness: int = -1,
+    color=(255, 0, 0),
+    kpt_thr: float = 0.3,
+    skeleton: list | None = None,
+    kpt_color=None,
+    link_color=None,
+    show_kpt_idx: bool = False,
+) -> np.ndarray:
+    import cv2
+    img = image.copy()
+    H, W = img.shape[:2]
+    if skeleton is None:
+        skeleton = []
+    if kpt_color is None:
+        kpt_color = color
+    if link_color is None:
+        link_color = (0, 255, 0)
+
+    def _as_color_list(c, n):
+        if hasattr(c, "detach"):
+            c = c.detach().cpu().numpy()
+        if isinstance(c, np.ndarray):
+            if c.ndim == 2 and c.shape[1] == 3:
+                return [tuple(int(v) for v in row) for row in c.tolist()]
+            if c.size == 3:
+                return [tuple(int(v) for v in c.tolist())] * max(1, n)
+        if isinstance(c, (list, tuple)):
+            if n and len(c) == n and isinstance(c[0], (list, tuple, np.ndarray)):
+                out = []
+                for cc in c:
+                    cc = np.asarray(cc).reshape(-1)
+                    out.append(tuple(int(v) for v in cc.tolist()))
+                return out
+            c_arr = np.asarray(c).reshape(-1)
+            if c_arr.size == 3:
+                return [tuple(int(v) for v in c_arr.tolist())] * max(1, n)
+        return [(255, 0, 0)] * max(1, n)
+
+    J = keypoints[0].shape[0] if keypoints else 0
+    kpt_colors = _as_color_list(kpt_color, J)
+    link_colors = _as_color_list(link_color, len(skeleton))
+
+    def in_bounds(x, y):
+        return 0 <= x < W and 0 <= y < H
+
+    for kpts, vis, score in zip(keypoints, keypoints_visible, keypoint_scores):
+        kpts = np.asarray(kpts, float)
+        vis = np.asarray(vis).reshape(-1).astype(bool)
+        score = np.asarray(score).reshape(-1)
+        for lk, (i, j) in enumerate(skeleton):
+            if i >= len(kpts) or j >= len(kpts):
+                continue
+            if not (vis[i] and vis[j]):
+                continue
+            if score[i] < kpt_thr or score[j] < kpt_thr:
+                continue
+            x1, y1 = map(int, np.round(kpts[i]))
+            x2, y2 = map(int, np.round(kpts[j]))
+            if not (in_bounds(x1, y1) and in_bounds(x2, y2)):
+                continue
+            cv2.line(img, (x1, y1), (x2, y2), link_colors[lk % len(link_colors)],
+                     thickness=max(1, thickness), lineType=cv2.LINE_AA)
+        for j_idx, (xy, v, s) in enumerate(zip(kpts, vis, score)):
+            if not v or s < kpt_thr:
+                continue
+            x, y = map(int, np.round(xy))
+            if not in_bounds(x, y):
+                continue
+            c = kpt_colors[min(j_idx, len(kpt_colors) - 1)]
+            cv2.circle(img, (x, y), radius, c, thickness=-1, lineType=cv2.LINE_AA)
+            if show_kpt_idx:
+                cv2.putText(img, str(j_idx), (x + radius, y - radius),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.4, c, 1, cv2.LINE_AA)
+    return img
+
+
+def _detect_persons(image_rgb: np.ndarray, threshold: float = 0.5):
+    import torch
+    proc, det = _get_detector()
+    pil_img = Image.fromarray(image_rgb)
+    inputs = proc(images=pil_img, return_tensors="pt")
+    with torch.no_grad():
+        outputs = det(**inputs)
+    target_sizes = torch.tensor([image_rgb.shape[:2]])
+    results = proc.post_process_object_detection(
+        outputs, target_sizes=target_sizes, threshold=threshold
+    )[0]
+    person_mask = results["labels"] == 1  # COCO class 1 = person
+    boxes = results["boxes"][person_mask].cpu().numpy()
+    scores = results["scores"][person_mask].cpu().numpy().reshape(-1, 1)
+    if len(boxes) == 0:
+        h, w = image_rgb.shape[:2]
+        return np.array([[0, 0, w - 1, h - 1, 1.0]], dtype=np.float32)
+    return np.concatenate([boxes, scores], axis=1).astype(np.float32)
+
+
+def _infer_pose(image_bgr, model, kpt_thr: float = 0.3):
+    import torch
+    import cv2
+
+    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+    bboxes = _detect_persons(image_rgb)
+    inputs_list, samples_list = [], []
+    for bbox in bboxes:
+        data_info = dict(img=image_bgr, bbox=bbox[None, :4], bbox_score=np.ones(1, dtype=np.float32))
+        data = model.pipeline(data_info)
+        data = model.data_preprocessor(data)
+        inputs_list.append(data["inputs"])
+        samples_list.append(data["data_samples"])
+    inputs = torch.cat(inputs_list, dim=0)
+    with torch.no_grad():
+        pred = model(inputs).cpu().numpy()
+
+    keypoints, scores = [], []
+    for i, sample in enumerate(samples_list):
+        kpts_i, scr_i = model.codec.decode(pred[i])
+        meta = sample["meta"] if isinstance(sample, dict) else sample.metainfo
+        kpts_i = kpts_i / np.array(meta["input_size"]) * meta["bbox_scale"] + meta["bbox_center"] - 0.5 * meta["bbox_scale"]
+        keypoints.append(kpts_i[0])
+        scores.append(scr_i[0])
+
+    pmeta = model.pose_metainfo
+    vis_rgb = image_rgb.copy()
+    # Scale render thickness so 308-keypoint dense pose stays visible on high-res input
+    short_side = min(vis_rgb.shape[:2])
+    radius_px = max(3, short_side // 200)
+    thick_px = max(2, short_side // 250)
+    box_thick = max(2, short_side // 300)
+    for bbox, kpts, scr in zip(bboxes, keypoints, scores):
+        x1, y1, x2, y2 = map(int, bbox[:4])
+        cv2.rectangle(vis_rgb, (x1, y1), (x2, y2), (0, 255, 0), box_thick)
+        vis_rgb = visualize_keypoints(
+            image=vis_rgb,
+            keypoints=[kpts],
+            keypoints_visible=[np.ones(len(scr), dtype=bool)],
+            keypoint_scores=[scr],
+            radius=radius_px, thickness=thick_px, kpt_thr=kpt_thr,
+            skeleton=pmeta["skeleton_links"],
+            kpt_color=pmeta["keypoint_colors"],
+            link_color=pmeta["skeleton_link_colors"],
+        )
+    return Image.fromarray(vis_rgb), f"persons={len(bboxes)} | kpts/person={len(keypoints[0]) if keypoints else 0}"
+
+
+# --- Predict entry point ----------------------------------------------------
+def predict(image: Image.Image, task: str, size: str):
+    if image is None:
+        return None, "No image provided"
+    key = (task, size)
+    if key not in VARIANTS:
+        return None, f"Unknown variant {task}-{size}. Allowed: {sorted(VARIANTS.keys())}"
+    t0 = time.time()
+    try:
+        import cv2
+        image_pil = image.convert("RGB")
+        in_w, in_h = image_pil.size
+        image_bgr = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
+        kind = VARIANTS[key]["kind"]
+        if size == "5b":
+            out_img, info = _infer_dense_5b(image_bgr, task)
+        elif kind == "pose":
+            model = _get_pose_model(size)
+            out_img, info = _infer_pose(image_bgr, model)
+        else:
+            model = _get_dense_model(task, size)
+            if kind == "seg":
+                out_img, info = _infer_seg(image_bgr, model)
+            elif kind == "normal":
+                out_img, info = _infer_normal(image_bgr, model)
+            elif kind == "pointmap":
+                out_img, info = _infer_pointmap(image_bgr, model)
+            else:
+                return None, f"Unhandled kind: {kind}"
+        elapsed = time.time() - t0
+        out_w, out_h = out_img.size
+        return out_img, f"{task}-{size}: done in {elapsed:.1f}s | {in_w}×{in_h} → 1024×768 → {out_w}×{out_h} | {info}"
+    except Exception as e:
+        return None, f"{type(e).__name__}: {e}\n\n{traceback.format_exc()[:1500]}"
+
+
+def health():
+    return (
+        f"Service up | dense cache: {list(_MODELS.keys())} | pose cache: {list(_POSE_MODELS.keys())} | "
+        f"detector_loaded={_DETECTOR is not None} | variants={len(VARIANTS)} "
+        f"({sorted(set(t for t, _ in VARIANTS))} × {sorted(set(s for _, s in VARIANTS))})"
+    )
+
+
+DEMO_IMAGES = sorted(str(p) for p in Path("/app/assets/images").glob("*.jpg"))
+
+with gr.Blocks(title="Sapiens2 CPU", css="""
+#img-in,#img-out{max-height:220px}
+#status-box textarea{max-height:60px!important;min-height:60px!important}
+#status-box{flex-grow:0!important}
+""") as demo:
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=1):
+            img_in = gr.Image(type="pil", label="Input", height=200, elem_id="img-in")
+            with gr.Row():
+                task_in = gr.Dropdown(choices=["seg", "normal", "pointmap", "pose"], value="seg", label="Task", scale=1)
+                size_in = gr.Dropdown(choices=["0.4b", "0.8b", "1b", "5b"], value="0.4b", label="Size", scale=1)
+            run_btn = gr.Button("Predict - 1024×768 native", variant="primary")
+            gr.Examples(
+                examples=[[u] for u in DEMO_IMAGES],
+                inputs=[img_in],
+                examples_per_page=6,
+                cache_examples=False,
+                label="Meta demo images",
+            )
+        with gr.Column(scale=1):
+            img_out = gr.Image(type="pil", label="Output", height=200, elem_id="img-out")
+            status = gr.Textbox(show_label=False, lines=2, max_lines=2, interactive=False, container=False, placeholder="Status will show here after Predict", elem_id="status-box")
+    run_btn.click(
+        fn=predict, inputs=[img_in, task_in, size_in], outputs=[img_out, status], api_name="predict"
+    )
+    # Keep health endpoint accessible via API (no UI button — useless in browser)
+    gr.Button(visible=False).click(fn=health, outputs=[gr.Textbox(visible=False)], api_name="health")
+
+demo.queue(default_concurrency_limit=1)
+
+if __name__ == "__main__":
+    demo.launch()

packages.txt

@@ -0,0 +1,2 @@
+libgl1
+libglib2.0-0

requirements.txt

@@ -0,0 +1,11 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch>=2.7
+torchvision
+pillow
+numpy
+safetensors
+huggingface_hub
+opencv-python-headless
+onnxruntime
+transformers
+sapiens @ git+https://github.com/facebookresearch/sapiens2.git