"""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])