| |
| """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) |
|
|
|
|
| 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: |
| x = image[0].float().div(255.0).permute(2, 0, 1) |
| else: |
| 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 |
| 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] |
| 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))) |
|
|
|
|
| |
|
|
| 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 |
|
|
|
|
| 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 |
|
|
|
|
| |
|
|
| 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)) |
|
|
|
|
| |
|
|
| @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() |
| with pytest.raises(RuntimeError): |
| m(_image()) |
|
|
|
|
| |
|
|
| 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 |
| 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): |
| 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 |
| 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() |
| assert fake.calls == calls |
| assert list(tmp_path.iterdir()) == [] |
| 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") |
| 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") |
|
|
