# SPDX-License-Identifier: Apache-2.0 """Device tests of the Python API (``tt_superpoint.SuperPoint``). One chip, opened by the API. cd code && python -m pytest -s -q models/tests/test_api_device.py 1. ``test_api_equals_server_and_timing``: ``model(...)`` returns the same keypoints, scores and descriptors as the HTTP server's request path (``models.server.app.predict_dict`` / ``_predict_core``) on the demo JPEG, for 7 parameter sets and 6 input types; then the warm ``model(...)`` time against the server path in the same process and the published numbers. 2. ``test_api_on_external_device_equals_server_pipeline``: the server's own model build and warm-up (``TtSuperPoint`` + ``_warmup_fused``) on a device opened by the caller, then ``SuperPoint.from_pretrained(device=...)`` on the same device: same output on the demo JPEG; ``close()`` leaves the caller's device open. """ from __future__ import annotations import base64 import io import os import statistics import time from pathlib import Path import numpy as np import pytest import torch from PIL import Image from tt_superpoint import SuperPoint CODE = Path(__file__).resolve().parents[2] SAMPLE = CODE / "sample_data" / "house_in_field_1080p.jpg" DEVICE_ID = int(os.environ.get("TT_DEVICE_ID", "0")) N_ITER = int(os.environ.get("SP_N_ITER", "200")) PARAM_SETS = [ {}, {"max_keypoints": -1}, {"max_keypoints": 300, "keypoint_threshold": 0.015}, {"nms_radius": 3}, # precompiled per-radius device NMS variant {"nms_radius": 0}, # host NMS fallback {"nms_radius": 12, "max_keypoints": 2000}, # host NMS fallback, > 1024 keypoints {"return_descriptors": False}, ] #: Published (OPT_REPORT.md round 10 / audit, ETH 12x10): served-like device_forward + postprocess #: on the decoded 1600x900 JPEG (the rgb case adds the host R-channel copy of the HWC array), and #: e2e_kpc_u8 on a 480x640 plane (H2D + trace + D2H + host decode; without the host sort by score). PUBLISHED_MS = {"plane_1600x900": (1.58 + 0.15, 1.68 + 0.17), "rgb_1600x900": (1.58 + 0.15, 1.68 + 0.17), "plane_480x640": (0.90, 0.96)} def _same(out, resp, params): assert resp["num_keypoints"] == len(out), (params, resp["num_keypoints"], len(out)) assert resp["keypoints"] == [[round(x, 3), round(y, 3)] for x, y in out.keypoints.double().tolist()], params assert resp["scores"] == [round(s, 6) for s in out.scores.tolist()], params if params.get("return_descriptors", True): d = np.load(io.BytesIO(base64.b64decode(resp["descriptors"]["data"])))["descriptors"] assert np.array_equal(d.view(np.uint16), out.descriptors.half().numpy().view(np.uint16)), params else: assert out.descriptors is None and "descriptors" not in resp def _attach_server(A, model): """Point the server module at the API's device model (the server's own request path runs).""" A.STATE.update(ready=True, model=model._m, tt_in=model._tt_in, fused=True, cfg={"fused": True}, model_config={"border_removal_distance": model.border_removal_distance}) def _detach_server(A): A.STATE.clear() A.STATE["ready"] = False def _time(f, n): ts = [] for _ in range(n): t0 = time.perf_counter() f() ts.append((time.perf_counter() - t0) * 1e3) return ts def test_api_equals_server_and_timing(): from models.server import app as A b64 = base64.b64encode(SAMPLE.read_bytes()).decode() with Image.open(SAMPLE) as im: im.load() rgb_pil = im.convert("RGB") rgb = np.asarray(rgb_pil) plane = np.ascontiguousarray(rgb[..., 0]) with SuperPoint.from_pretrained(device_id=DEVICE_ID, precompile_nms_radii=(3,)) as model: print(f"\nmodel: {model!r} config={model.config}") _attach_server(A, model) try: # 1) equality with the server response, every parameter set for params in PARAM_SETS: resp = A.predict_dict(A.PredictRequest(image=b64, **params)) out = model(str(SAMPLE), **params) _same(out, resp, params) exp_dev_nms = params.get("nms_radius", 4) in range(1, 9) assert out.device_nms == exp_dev_nms == resp["serving_path"]["device_nms"], params print(f"same as server: params={params} n={len(out)} device_nms={out.device_nms}") # 2) every input type -> the same output as the file path ref = model(str(SAMPLE)) inputs = { "bytes": SAMPLE.read_bytes(), "PIL": rgb_pil, "numpy HWC RGB": rgb, "numpy HWC BGR": (rgb[..., ::-1].copy(), {"bgr": True}), "numpy plane": plane, "torch CHW float": torch.from_numpy(rgb.copy()).permute(2, 0, 1).float() / 255.0, } for name, x in inputs.items(): x, kw = x if isinstance(x, tuple) else (x, {}) o = model(x, **kw) assert torch.equal(o.keypoints, ref.keypoints) and torch.equal(o.scores, ref.scores), name assert torch.equal(o.descriptors, ref.descriptors), name outs = model([str(SAMPLE), rgb, plane]) assert all(torch.equal(o.descriptors, ref.descriptors) for o in outs) # the 480x640 plane: the server's /predict_plane path p480 = np.array(Image.fromarray(plane).resize((640, 480), Image.BILINEAR)) resp = A._predict_core(A.PlaneParams(), None, p480) _same(model(p480), resp, {}) print(f"input types identical: {list(inputs)} + list call + 480x640 plane") # 3) warm timing, model(...) vs the server's request path in the same process (alternating) pp = A.PlaneParams() cases = { "plane_1600x900": (lambda: model(plane), lambda: A._predict_core(pp, None, plane)), "rgb_1600x900": (lambda: model(rgb), lambda: A._predict_core(pp, None, plane)), "plane_480x640": (lambda: model(p480), lambda: A._predict_core(pp, None, p480)), "jpeg_path": (lambda: model(str(SAMPLE)), lambda: A.predict_dict(A.PredictRequest(image=b64))), } rows = {} for name, (fa, fs) in cases.items(): for f in (fa, fs): _time(f, 20) api, srv = [], [] for _ in range(4): api += _time(fa, N_ITER // 4) srv += _time(fs, N_ITER // 4) rows[name] = (statistics.median(api), min(api), statistics.median(srv), min(srv)) pub = PUBLISHED_MS.get(name) print(f"timing {name}: model() median {rows[name][0]:.3f} min {rows[name][1]:.3f} ms | " f"server path median {rows[name][2]:.3f} min {rows[name][3]:.3f} ms" + (f" | published {pub[0]:.2f}-{pub[1]:.2f} ms" if pub else "")) for name, (am, amin, sm, smin) in rows.items(): # same request path: not slower than the server path measured alongside (5 % + 50 us noise) assert am <= sm * 1.05 + 0.05, (name, am, sm) finally: _detach_server(A) def test_api_on_external_device_equals_server_pipeline(): import ttnn from models.reference.superpoint_reference import load_reference_model from models.server import app as A from models.tt.superpoint_ttnn import TtSuperPoint from tt_superpoint import DEFAULT_REVISION from tt_superpoint.device import open_device b64 = base64.b64encode(SAMPLE.read_bytes()).decode() dev, mode = open_device(DEVICE_ID) try: # the server's own build + warm-up (models/server/app.py lifespan, fused path) tm = load_reference_model(revision=DEFAULT_REVISION) torch.set_grad_enabled(False) sp = TtSuperPoint(tm, dev, input_height=A.INPUT_HEIGHT, input_width=A.INPUT_WIDTH, fused=True) tt_in = sp.allocate_input(batch_size=1) A.STATE.update(model=sp, tt_in=tt_in, fused=True, cfg={"fused": True}, model_config={"border_removal_distance": int(tm.config.border_removal_distance)}) A._warmup_fused(sp, tt_in, torch.zeros(1, 3, A.INPUT_HEIGHT, A.INPUT_WIDTH), ttnn, dev) A.STATE["ready"] = True resps = [A.predict_dict(A.PredictRequest(image=b64, **p)) for p in PARAM_SETS[:3]] _detach_server(A) ttnn.synchronize_device(dev) sp.release() ttnn.deallocate(tt_in) del sp with SuperPoint.from_pretrained(device=dev) as model: for p, resp in zip(PARAM_SETS[:3], resps): _same(model(str(SAMPLE), **p), resp, p) print(f"\nexternal device ({mode}): API output == server pipeline output for {PARAM_SETS[:3]}") g = dev.compute_with_storage_grid_size() # still open after close() assert (g.x, g.y) == (12, 10) or mode == "worker" finally: _detach_server(A) ttnn.close_device(dev)