inference.py

"""Tactus inference: open-vocabulary object recognition from pressure-array data.

Tactus embeds a short window of 32x32 tactile pressure frames into the fusion-embedding
canonical 2048-d text space, so recognition is a cosine ranking against text queries rather
than a fixed classifier head.

Two halves, loaded from two places:
  * PRESSURE side (this repository): a trained CNN trunk + projector (``tactus_head.pt``),
    defined in the vendored ``tactile.py`` (byte-identical to the training-time module, so the
    checkpoint loads with ``strict=True``).
  * TEXT side (the base model repository): the canonical whitened text readout of
    ``EximiusLabs/fusion-embedding-2-2b-preview`` via the ``fusion-embedding`` package.
    Tactus was trained against that exact readout; embedding text any other way (including the
    plain base model without whitening) will misrank.

Input contract: frames are STAG-normalized pressure maps. If your data is raw sensor counts,
convert with ``stag_raw_to_unit`` (the dataset's own calibration: ``clip((raw-500)/150, 0, 1)``)
before embedding, or pass ``raw="stag"``. uint8 0-255 maps produced by the STAG caches are
accepted directly.

Usage:
    from inference import TactusEmbedder

    ta = TactusEmbedder.from_pretrained("EximiusLabs/fusion-embedding-2-tactus",
                                        revision="v0.1-preview")
    window = np.load("grasp.npy")            # [F,32,32] uint8/float, or [32,32]
    for text, score in ta.rank(window, ["a mug", "scissors", "a full soda can"]):
        print(f"{score:+.3f}  {text}")
"""
from __future__ import annotations

import os

import numpy as np
import torch
import torch.nn.functional as F

from tactile import TactileEncoder, build_projector, preprocess_pressure  # vendored, verbatim

REPO_DEFAULT = "EximiusLabs/fusion-embedding-2-tactus"
BASE_REPO = "EximiusLabs/fusion-embedding-2-2b-preview"
SAFETENSORS_FILE = "model.safetensors"      # head weights; config in the file's metadata
CKPT_FILE = "tactus_head.pt"                # legacy .pt with full training provenance

# STAG's own sensor calibration (classification/TouchDataset.py::transformPressure in the
# reference code): resting level raw ~500 counts, informative band [500, 650].
STAG_RAW_BASELINE = 500.0
STAG_RAW_SCALE = 150.0


def stag_raw_to_unit(raw) -> np.ndarray:
    """Raw pressure counts -> [0,1] floats via STAG's calibration affine."""
    a = np.asarray(raw, dtype=np.float32)
    return np.clip((a - STAG_RAW_BASELINE) / STAG_RAW_SCALE, 0.0, 1.0)


class TactusEmbedder:
    """Pressure -> 2048-d canonical embeddings, plus the matching text side."""

    def __init__(self, head_path: str, device: str = "cuda", dtype=torch.bfloat16,
                 base_repo: str = BASE_REPO, load_text: bool = True):
        self.device = device
        if head_path.endswith(".safetensors"):
            import json as _json

            import safetensors.torch as st
            from safetensors import safe_open
            flat = st.load_file(head_path)
            with safe_open(head_path, framework="pt") as f:
                cfg = _json.loads((f.metadata() or {})["config"])
            enc_sd = {k[len("encoder."):]: v for k, v in flat.items() if k.startswith("encoder.")}
            proj_sd = {k[len("proj."):]: v for k, v in flat.items() if k.startswith("proj.")}
        else:                                          # legacy .pt checkpoint
            blob = torch.load(head_path, map_location="cpu", weights_only=False)
            cfg = blob["config"]
            enc_sd = {k: v.float() for k, v in blob["encoder"].items()}
            proj_sd = {k: v.float() for k, v in blob["proj"].items()}
        self.cfg = cfg
        self.window_frames = int(cfg.get("window_frames", 8))

        self.enc = TactileEncoder(temporal=cfg.get("temporal", "conv"),
                                  depth=cfg.get("depth", "resnet18"),
                                  frames=self.window_frames,
                                  flow=bool(cfg.get("flow", False)))
        self.proj = build_projector()
        self.enc.load_state_dict(enc_sd, strict=True)
        self.proj.load_state_dict(proj_sd, strict=True)
        self.enc.eval().to(device)
        self.proj.eval().to(device)

        # The text side is the canonical whitened readout Tactus was trained against.
        self.ue = None
        if load_text:
            from fusion_embedding import UnifiedEmbedder   # pip install fusion-embedding[hf]
            self.ue = UnifiedEmbedder.from_pretrained(base_repo, device=device, dtype=dtype)

    @classmethod
    def from_pretrained(cls, repo: str = REPO_DEFAULT, revision: str | None = None,
                        device: str = "cuda", dtype=torch.bfloat16, load_text: bool = True):
        if os.path.isdir(repo):
            head = None
            for name in (SAFETENSORS_FILE, CKPT_FILE):
                for cand in (os.path.join(repo, name), os.path.join(repo, "out", name)):
                    if os.path.exists(cand):
                        head = cand
                        break
                if head:
                    break
            assert head, f"no {SAFETENSORS_FILE} or {CKPT_FILE} under {repo}"
        else:
            from huggingface_hub import hf_hub_download
            try:
                head = hf_hub_download(repo, SAFETENSORS_FILE, revision=revision)
            except Exception:  # noqa: BLE001 -- older revisions ship only the .pt
                head = hf_hub_download(repo, CKPT_FILE, revision=revision)
        return cls(head, device=device, dtype=dtype, load_text=load_text)

    @torch.no_grad()
    def embed_pressure(self, frames, raw: str | None = None) -> torch.Tensor:
        """One grasp window ``[F,32,32]`` (or a single ``[32,32]`` frame) -> 2048-d, L2-normalized.

        ``raw="stag"`` applies the sensor calibration affine first (input is raw counts).
        Otherwise frames must already be [0,1] floats or uint8 0-255 maps (the cache format)."""
        if raw == "stag":
            frames = stag_raw_to_unit(frames)
        elif raw is not None:
            raise ValueError(f"unknown raw mode {raw!r}; use raw='stag' or pre-normalize")
        x = preprocess_pressure(frames).to(self.device)        # [F,32,32] in [0,1]
        feat = self.enc(x)                                     # [512]
        vec = self.proj(feat)                                  # [2048] raw
        return F.normalize(vec.float(), dim=-1).cpu()

    @torch.no_grad()
    def embed_text(self, texts) -> torch.Tensor:
        """Text -> canonical 2048-d embeddings (whitened readout). Requires load_text=True."""
        assert self.ue is not None, "constructed with load_text=False"
        if isinstance(texts, str):
            texts = [texts]
        return torch.stack([self.ue.embed_text(t).float() for t in texts])

    @torch.no_grad()
    def rank(self, frames, texts, raw: str | None = None):
        """Rank candidate texts against one pressure window; returns [(text, cosine)] sorted."""
        p = self.embed_pressure(frames, raw=raw)               # [2048]
        t = F.normalize(self.embed_text(texts), dim=-1)        # [N,2048]
        sims = (t @ p).tolist()
        return sorted(zip(list(texts), sims), key=lambda kv: -kv[1])