# SPDX-License-Identifier: Apache-2.0 """Host-only tests of the Python API (``tt_moge.MoGeModel``): no device, no ttnn, no weights. TT_VISIBLE_DEVICES=none python -m pytest tt_moge/tests/test_api.py -q A CPU stand-in for the device model (``_FakeTt``) produces canvas outputs that depend on the input pixels. The API's output must equal the HTTP server's ``/predict`` response computed from the same stand-in (``tt_moge.server.app.predict``, ``npz_binary``), for every request option. """ from __future__ import annotations import base64 import io import os import subprocess import sys import numpy as np import pytest import torch from PIL import Image import tt_moge from tt_moge import MoGeModel, MoGeOutput from tt_moge.api import _load_rgb CANVAS = (320, 240) # small canvas (>= 224) keeps the host test fast class _FakeTt: """Stands in for tt_moge.tt.ttnn_moge.TtMoGe: input canvas -> canvas-resolution outputs.""" def __init__(self): self.calls = 0 self.released = False def __call__(self, image, num_tokens): self.calls += 1 if image.dtype == torch.uint8: # (1, H, W, 3) x = image[0].float().div(255.0).permute(2, 0, 1) else: # (1, 3, H, W) x = image[0].float() _, h, w = x.shape ar = w / h u = torch.linspace(-ar / (1 + ar ** 2) ** 0.5, ar / (1 + ar ** 2) ** 0.5, w) v = torch.linspace(-1 / (1 + ar ** 2) ** 0.5, 1 / (1 + ar ** 2) ** 0.5, h) vv, uu = torch.meshgrid(v, u, indexing="ij") z = 1.0 + x.mean(0) + 0.3 * vv # depends on the pixels f = 1.3 points = torch.stack([uu * (z + 0.2) / f, vv * (z + 0.2) / f, z], dim=-1)[None] n = torch.stack([0.1 * uu, 0.1 * vv - 1.0, -torch.ones_like(z)], dim=-1) normal = torch.nn.functional.normalize(n, dim=-1)[None] mask = torch.sigmoid(4.0 * (x[0] - 0.05))[None] # dark pixels invalid return {"points": points.contiguous(), "normal": normal.contiguous(), "mask": mask.contiguous(), "metric_scale": torch.tensor([1.7])} def release(self): self.released = True def _image(w=200, h=150, seed=0) -> np.ndarray: rng = np.random.default_rng(seed) yy, xx = np.mgrid[0:h, 0:w] base = np.stack([(xx * 255 // w), (yy * 255 // h), ((xx + yy) * 127 // (w + h))], -1) return np.clip(base + rng.integers(0, 40, (h, w, 3)), 0, 255).astype(np.uint8) @pytest.fixture() def model(): m = MoGeModel(_FakeTt(), device=None, resolution=CANVAS, num_tokens=1800) yield m m.close() def _server_npz(img: np.ndarray, monkeypatch, **req) -> dict: """The server's /predict on ``img`` with the same stand-in model (npz_binary body).""" pytest.importorskip("fastapi", reason="the server comparison needs the [server] extra") monkeypatch.setenv("MOGE_CANONICAL_SIZE", f"{CANVAS[0]}x{CANVAS[1]}") monkeypatch.setenv("MOGE_NUM_TOKENS", "1800") from tt_moge.server import app, fastpath buf = io.BytesIO() Image.fromarray(img).save(buf, format="PNG") saved = dict(app.STATE) try: app.STATE.clear() app.STATE.update(config=app.Config(), model=_FakeTt(), fast=fastpath.enabled()) resp = app.predict(app.PredictRequest(image=base64.b64encode(buf.getvalue()).decode(), output_format="npz_binary", **req)) finally: app.STATE.clear() app.STATE.update(saved) body = bytes(resp.body) return dict(np.load(io.BytesIO(body))) # ------------------------------------------------------------------ package surface def test_import_is_light(): code = ("import sys, tt_moge; from tt_moge import MoGeModel, MoGeOutput; " "assert 'ttnn' not in sys.modules, 'importing the API must not import ttnn'; print('ok')") r = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, env={**os.environ, "TT_VISIBLE_DEVICES": "none"}) assert r.returncode == 0, r.stderr assert "ok" in r.stdout def test_public_names(): assert tt_moge.__version__ for name in ("from_pretrained", "__call__", "predict", "predict_many", "infer", "forward", "close", "__enter__", "__exit__", "warmup", "warmed_variants"): assert hasattr(MoGeModel, name), name with pytest.raises(AttributeError): tt_moge.NotAName # noqa: B018 def test_from_pretrained_validates_before_device(): with pytest.raises(ValueError): MoGeModel.from_pretrained(num_tokens=100) with pytest.raises(ValueError): MoGeModel.from_pretrained(resolution=(100, 100)) def test_no_machine_paths_in_runtime_code(): root = os.path.dirname(tt_moge.__file__) files = [os.path.join(root, "api.py"), os.path.join(root, "device.py"), os.path.join(root, "tt", "mm2d.py"), os.path.join(root, "tt", "host_ext.py")] files += [os.path.join(root, "server", f) for f in ("app.py", "fastpath.py", "postprocess.py")] for f in files: src = open(f).read() assert "/home/" not in src, f # ------------------------------------------------------------------ input conversion def test_load_rgb_variants(tmp_path): a = _image() p = tmp_path / "x.png" Image.fromarray(a).save(p) for x in (str(p), p, Image.fromarray(a), a, torch.from_numpy(a), Image.fromarray(a).convert("RGBA"), np.concatenate([a, np.full(a.shape[:2] + (1,), 255, np.uint8)], -1)): r = _load_rgb(x) assert isinstance(r, np.ndarray) and r.dtype == np.uint8 np.testing.assert_array_equal(r, a) g = _load_rgb(a[..., 0]) assert g.shape == a.shape and (g[..., 1] == a[..., 0]).all() f = a.astype(np.float32) / 255 for x in (f, np.moveaxis(f, -1, 0), torch.from_numpy(f).permute(2, 0, 1), torch.from_numpy(f).permute(2, 0, 1)[None]): r = _load_rgb(x) assert isinstance(r, torch.Tensor) and r.shape == (1, 3, a.shape[0], a.shape[1]) assert torch.equal(r[0].permute(1, 2, 0), torch.from_numpy(f)) with pytest.raises(TypeError): _load_rgb(a.astype(np.int32)) with pytest.raises(TypeError): _load_rgb(12) with pytest.raises(ValueError): _load_rgb(torch.zeros(2, 3, 8, 8)) # ------------------------------------------------------------------ the call vs the server @pytest.mark.parametrize("opts", [ {}, {"fit": "stretch"}, {"fov_x": 60.0}, {"apply_mask": False}, {"force_projection": False}, {"fit": "stretch", "fov_x": 75.0, "apply_mask": False}, ]) def test_call_matches_server(model, monkeypatch, opts): img = _image(w=200, h=150) out = model(img, **opts) assert isinstance(out, MoGeOutput) srv = _server_npz(img, monkeypatch, **opts) assert out.depth.shape == (150, 200) and out.points.shape == (150, 200, 3) and out.normal.shape == (150, 200, 3) assert out.depth.dtype == np.float32 and out.points.dtype == np.float32 and out.normal.dtype == np.float32 assert out.mask.dtype == np.bool_ and out.intrinsics.shape == (3, 3) 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"] assert 0.0 < out.mask.mean() <= 1.0 def test_output_container(model, tmp_path): out = model(_image()) assert set(out.keys()) == {"points", "depth", "normal", "mask", "intrinsics", "metric_scale"} assert out["depth"] is out.depth and "mask" in out and "nope" not in out k = out.intrinsics_pixels assert k[0, 0] == pytest.approx(out.intrinsics[0, 0] * out.width) assert k[1, 2] == pytest.approx(out.intrinsics[1, 2] * out.height) t = out.to_torch() assert set(t) == {"points", "intrinsics", "depth", "mask", "normal"} and t["mask"].dtype == torch.bool out.save_npz(tmp_path / "o.npz") z = np.load(tmp_path / "o.npz") assert z["normal"].dtype == np.float16 and z["mask"].dtype == np.uint8 assert "MoGeOutput(200x150" in repr(out) v = out.mask assert np.isfinite(out.depth[v]).all() and np.isinf(out.depth[~v]).all() def test_inputs_give_same_result(model, tmp_path): a = _image() p = tmp_path / "x.png" Image.fromarray(a).save(p) ref = model(a) for x in (str(p), Image.fromarray(a), torch.from_numpy(a)): o = model(x) np.testing.assert_array_equal(o.depth, ref.depth) np.testing.assert_array_equal(o.normal, ref.normal) def test_float_canvas_equals_uint8_canvas(model): a = _image(*CANVAS) f = torch.from_numpy(a).permute(2, 0, 1)[None].float() / 255 o8, of = model(a), model(f) np.testing.assert_array_equal(o8.depth, of.depth) np.testing.assert_array_equal(o8.points, of.points) def test_predict_many_matches_single_calls(model): imgs = [_image(seed=i, w=180 + 10 * i, h=140) for i in range(4)] many = model.predict_many(imgs, fov_x=70.0) as_list = model(imgs, fov_x=70.0) for img, a, b in zip(imgs, many, as_list): s = model(img, fov_x=70.0) for o in (a, b): np.testing.assert_array_equal(o.points, s.points) np.testing.assert_array_equal(o.mask, s.mask) assert o.metric_scale == s.metric_scale assert model.predict_many([]) == [] def test_infer_is_upstream_shaped(model): a = _image(w=160, h=120) x = torch.from_numpy(a).permute(2, 0, 1).float() / 255 d = model.infer(x) assert set(d) == {"points", "intrinsics", "depth", "mask", "normal"} assert d["points"].shape == (120, 160, 3) and d["depth"].shape == (120, 160) and d["intrinsics"].shape == (3, 3) b = model.infer(torch.stack([x, x])) assert b["points"].shape == (2, 120, 160, 3) and b["intrinsics"].shape == (2, 3, 3) assert torch.equal(b["depth"][1], d["depth"]) with pytest.raises(ValueError): model.infer(x, num_tokens=1200) def test_forward_requires_canvas(model): a = _image(*CANVAS) out = model.forward(a) assert out["points"].shape == (1, CANVAS[1], CANVAS[0], 3) with pytest.raises(ValueError): model.forward(_image(100, 100)) def test_argument_checks(model): with pytest.raises(ValueError): model(_image(), fit="crop") with pytest.raises(ValueError): model(_image(), fov_x=200) with pytest.raises(ValueError): model(np.zeros((4, 4, 3), np.uint8)) def test_close_and_context_manager(): fake = _FakeTt() with MoGeModel(fake, device=None, resolution=CANVAS) as m: m(_image()) assert "open" in repr(m) assert fake.released and m.closed m.close() # idempotent with pytest.raises(RuntimeError): m(_image()) # ------------------------------------------------------------------ warm-up API def test_warmup_spec_in_signature(): import inspect from tt_moge.api import DEFAULT_WARMUP_VARIANTS, WARMUP_PRESETS p = inspect.signature(MoGeModel.from_pretrained).parameters assert p["warmup_variants"].default == DEFAULT_WARMUP_VARIANTS == "default" assert p["verbose"].default is False assert {"none", "minimal", "default", "all"} <= set(WARMUP_PRESETS) assert all(v in WARMUP_PRESETS["all"] for v in WARMUP_PRESETS["default"]) @pytest.mark.parametrize("canvas", [CANVAS, (1920, 1080)]) def test_warmup_presets_are_valid(canvas): from tt_moge.api import WARMUP_PRESETS, _expand_warmup_spec, _normalize_variant for name in WARMUP_PRESETS: keys = [_normalize_variant(v, canvas) for v in _expand_warmup_spec(name)] assert len(keys) == len(set(keys)), name # no duplicate variant in a preset assert _expand_warmup_spec(None) == [] and _expand_warmup_spec("none") == [] assert _expand_warmup_spec(["minimal", {"fit": "stretch"}]) == [{}, {"fit": "stretch"}] @pytest.mark.parametrize("bad, exc", [ ("bogus", ValueError), ({"input": "gif"}, ValueError), ({"call": "forward", "size": (800, 600)}, ValueError), ({"size": (4, 4)}, ValueError), ({"fov_x": 500}, ValueError), ({"fit": "crop"}, ValueError), ({"colour": 1}, ValueError), (5, TypeError), ]) def test_from_pretrained_rejects_bad_warmup_before_device(bad, exc): with pytest.raises(exc): # raised before weights / device (none here) MoGeModel.from_pretrained(warmup_variants=bad) def test_warmup_runs_each_variant_once(model, tmp_path, monkeypatch): import tempfile monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) fake = model._tt n0 = fake.calls s = model.warmup(input="jpeg", size=(200, 150), fit="stretch") assert s > 0 and fake.calls == n0 + 1 assert model.warmup(input="jpeg", size=(200, 150), fit="stretch") == 0.0 # idempotent assert model.warmup(input="jpeg", size=(200, 150), fit="stretch", fov_x=None) == 0.0 assert fake.calls == n0 + 1 model.warmup(["all"]) from tt_moge.api import WARMUP_PRESETS n_all = len(WARMUP_PRESETS["all"]) rows = model.warmed_variants assert len(rows) == n_all + 1 and all(r["warmup_ms"] >= 0 for r in rows) assert {r["input"] for r in rows} == {"array", "float", "png", "jpeg", "pil"} calls = fake.calls model.warmup("all") model.warmup() # the default variant: already prepared assert fake.calls == calls assert list(tmp_path.iterdir()) == [] # the temporary png / jpeg files are removed with pytest.raises(TypeError): model.warmup(["minimal"], fit="pad") def test_warmup_does_not_change_outputs(tmp_path): img = _image(w=220, h=160, seed=3) p = tmp_path / "a.png" Image.fromarray(img).save(p) cold = MoGeModel(_FakeTt(), device=None, resolution=CANVAS) ref = [cold(img), cold(str(p), fov_x=70.0), cold(img, fit="stretch", apply_mask=False)] cold.close() warm = MoGeModel(_FakeTt(), device=None, resolution=CANVAS) warm.warmup("all") got = [warm(img), warm(str(p), fov_x=70.0), warm(img, fit="stretch", apply_mask=False)] warm.close() for a, b in zip(ref, got): for k in a.keys(): np.testing.assert_array_equal(np.asarray(a[k]), np.asarray(b[k])) def test_synthetic_warmup_image(): from tt_moge.api import _synthetic_image a = _synthetic_image(64, 48) assert a.shape == (48, 64, 3) and a.dtype == np.uint8 assert np.array_equal(a, _synthetic_image(64, 48)) and a.std() > 10 def test_quiet_logs_only_while_loading(tmp_path, monkeypatch): """verbose=False sets TT_LOGGER_LEVEL / LOGURU_LEVEL for the ttnn import and the device open only; a user's own value wins; verbose=True changes nothing.""" from tt_moge.api import _quiet_tt_logs pkg = tmp_path / "ttnn" pkg.mkdir() (pkg / "__init__.py").write_text("import os\nSEEN = (os.environ.get('TT_LOGGER_LEVEL'), os.environ.get('LOGURU_LEVEL'))\n") monkeypatch.syspath_prepend(str(tmp_path)) monkeypatch.delitem(sys.modules, "ttnn", raising=False) monkeypatch.delenv("TT_LOGGER_LEVEL", raising=False) monkeypatch.setenv("LOGURU_LEVEL", "DEBUG") # the user's choice with _quiet_tt_logs(verbose=True): assert "ttnn" not in sys.modules and "TT_LOGGER_LEVEL" not in os.environ with _quiet_tt_logs(verbose=False): assert os.environ["TT_LOGGER_LEVEL"] == "error" and os.environ["LOGURU_LEVEL"] == "DEBUG" assert sys.modules["ttnn"].SEEN == ("error", "DEBUG") assert "TT_LOGGER_LEVEL" not in os.environ and os.environ["LOGURU_LEVEL"] == "DEBUG" monkeypatch.delitem(sys.modules, "ttnn")