File size: 6,139 Bytes
ea621a3 2a74422 ea621a3 2a74422 ea621a3 2a74422 ea621a3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | # SPDX-License-Identifier: Apache-2.0
"""Device tests of the Python API: ``model(image)`` vs the HTTP server's /predict, and call latency.
python -m pytest tt_moge/tests/test_api_device.py -q -s # one chip; device_id from MOGE_DEVICE (0)
Run this file on its own (it opens the device through ``MoGeModel.from_pretrained``, the published
configuration: ETH dispatch, 1 command queue, 12x10 grid -- the p150 configuration), not in one session with the tests
that use the ``device`` fixture of conftest.py.
* equivalence: for the demo image (``media/source.png``, 1920x1080) and an 800x600 copy, every array of
``model(image, **opts)`` equals the server's ``npz`` response (``tt_moge.server.app.predict`` on the
same TtMoGe object), bit for bit (``normal`` compared after the server's float16 cast).
* speed: warm repeated calls; ``model.forward(image)`` is the published ``tt(image)`` call
(README / OPT_REPORT.md "p150 ETH-dispatch compliance 2026-10-05", chip 1: median 20.31-21.06 ms, min
19.24-19.49 ms). Its median must not be slower than 21.06 ms x (1 + MOGE_PERF_TOL), MOGE_PERF_TOL default 0.05
(shared host).
"""
from __future__ import annotations
import base64
import io
import json
import os
import statistics
import time
import numpy as np
import pytest
import torch
from PIL import Image
if os.environ.get("TT_VISIBLE_DEVICES", "").lower() == "none":
pytest.skip("device test (TT_VISIBLE_DEVICES=none)", allow_module_level=True)
ttnn = pytest.importorskip("ttnn")
from tt_moge import MoGeModel # noqa: E402
REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
DEMO = os.path.join(REPO, "media", "source.png")
PUBLISHED_CALL_MS = {"median": (20.31, 21.06), "min": (19.24, 19.49)} # tt(image), ETH + 1 CQ, chip 1, 2026-10-05
ITERS = int(os.environ.get("MOGE_PERF_ITERS", "60"))
@pytest.fixture(scope="module")
def model():
m = MoGeModel.from_pretrained(device_id=int(os.environ.get("MOGE_DEVICE", "0")), verbose=True)
yield m
m.close()
def _png_b64(img: np.ndarray) -> str:
buf = io.BytesIO()
Image.fromarray(img).save(buf, format="PNG")
return base64.b64encode(buf.getvalue()).decode()
def _server(model, img: np.ndarray, **opts) -> dict:
from tt_moge.server import app, fastpath
saved = dict(app.STATE)
try:
app.STATE.clear()
app.STATE.update(config=app.Config(), model=model._tt, fast=fastpath.enabled())
resp = app.predict(app.PredictRequest(image=_png_b64(img), output_format="npz_binary", **opts))
finally:
app.STATE.clear()
app.STATE.update(saved)
return dict(np.load(io.BytesIO(bytes(resp.body))))
def _assert_same(out, srv):
np.testing.assert_array_equal(out.depth, srv["depth"])
np.testing.assert_array_equal(out.points, srv["points"])
np.testing.assert_array_equal(out.normal.astype(np.float16), srv["normal"])
np.testing.assert_array_equal(out.mask.astype(np.uint8), srv["mask"])
np.testing.assert_array_equal(out.intrinsics, srv["intrinsics"])
assert np.float32(out.metric_scale) == srv["metric_scale"]
@pytest.mark.parametrize("opts", [{}, {"fit": "stretch", "fov_x": 60.0},
{"apply_mask": False, "force_projection": False}])
def test_demo_matches_server(model, opts):
img = np.asarray(Image.open(DEMO).convert("RGB"))
out = model(DEMO, **opts)
srv = _server(model, img, **opts)
_assert_same(out, srv)
print(f"\n[api] {opts or 'defaults'}: {out!r}")
def test_resized_matches_server(model):
img = np.asarray(Image.open(DEMO).convert("RGB").resize((800, 600), Image.BILINEAR))
out = model(img)
assert out.depth.shape == (600, 800)
_assert_same(out, _server(model, img))
def test_input_types_agree(model):
img = np.array(Image.open(DEMO).convert("RGB"))
ref = model(DEMO)
f = torch.from_numpy(img).permute(2, 0, 1).float() / 255 # upstream-style float (3, H, W)
for x in (img, Image.open(DEMO), f, f[None]):
o = model(x)
np.testing.assert_array_equal(o.points, ref.points)
np.testing.assert_array_equal(o.normal, ref.normal)
many = model.predict_many([DEMO, img])
for o in many:
np.testing.assert_array_equal(o.depth, ref.depth)
d = model.infer(f)
assert torch.equal(d["depth"], torch.from_numpy(ref.depth))
def _times(fn, n):
ts = []
for _ in range(n):
t0 = time.perf_counter()
fn()
ts.append(1e3 * (time.perf_counter() - t0))
return statistics.median(ts), min(ts)
def test_call_latency(model):
from tt_moge.benchmark import load_image
img_f = load_image(DEMO).float() # the bench's input: float32 (1, 3, 1080, 1920)
img_u8 = np.asarray(Image.open(DEMO).convert("RGB"))
for _ in range(5):
model.forward(img_f)
model(img_u8)
res = {
"forward_float": _times(lambda: model.forward(img_f), ITERS), # == tt(image) of bench_breakdown e2e
"forward_u8": _times(lambda: model.forward(img_u8), ITERS), # the served u8 canvas input
"call_u8_array": _times(lambda: model(img_u8), ITERS), # + focal/shift solve + finish
"call_path": _times(lambda: model(DEMO), ITERS // 2), # + PNG decode
}
n = 8
t0 = time.perf_counter()
model.predict_many([img_u8] * n)
res["predict_many_per_image"] = (1e3 * (time.perf_counter() - t0) / n,) * 2
for k, (med, mn) in res.items():
print(f"\n[api perf] {k:24s} median {med:7.2f} ms min {mn:7.2f} ms")
print(f"[api perf] published tt(image): median {PUBLISHED_CALL_MS['median']} ms, min {PUBLISHED_CALL_MS['min']} ms")
if os.environ.get("MOGE_API_PERF_JSON"):
with open(os.environ["MOGE_API_PERF_JSON"], "w") as f:
json.dump({k: {"median_ms": v[0], "min_ms": v[1]} for k, v in res.items()}
| {"published_tt_image_ms": PUBLISHED_CALL_MS, "iters": ITERS}, f, indent=1)
tol = float(os.environ.get("MOGE_PERF_TOL", "0.05"))
assert res["forward_float"][0] <= PUBLISHED_CALL_MS["median"][1] * (1 + tol), res
|