| |
| """Device test: after ``MoGeModel.from_pretrained`` (default ``warmup_variants``), the FIRST real call of |
| every default warm-up variant is as fast as the later calls. |
| |
| python -m pytest tt_moge/tests/test_api_first_call.py -q -s # one chip; device_id from MOGE_DEVICE (0) |
| |
| Run this file on its own, in a fresh process (it opens the device through ``from_pretrained``). Paths |
| are resolved relative to this file, so the working directory does not matter. |
| |
| For each variant, the inputs are NEW real images (crops / flips of ``media/source.png``, never the |
| warm-up image): the first call must be within ``MOGE_FIRST_CALL_TOL`` (default 0.10) of the median of |
| the next 10 calls, plus 1.5 ms for the jitter of a shared host. The first call's output must equal a |
| repeat call on the same input (warm-up keeps no state that changes results). |
| """ |
| from __future__ import annotations |
|
|
| import os |
| import statistics |
| import time |
| import zlib |
|
|
| 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) |
| pytest.importorskip("ttnn") |
|
|
| from tt_moge import MoGeModel |
|
|
| REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) |
| DEMO = os.path.join(REPO, "media", "source.png") |
| TOL = float(os.environ.get("MOGE_FIRST_CALL_TOL", "0.10")) |
| SLACK_MS = 1.5 |
| N_STEADY = 10 |
|
|
|
|
| @pytest.fixture(scope="module") |
| def model(): |
| t = time.perf_counter() |
| m = MoGeModel.from_pretrained(device_id=int(os.environ.get("MOGE_DEVICE", "0"))) |
| print(f"\n[first-call] from_pretrained {time.perf_counter() - t:.1f} s; warmed: " |
| + ", ".join(MoGeModel._variant_str(tuple(r[k] for k in ("input", "size", "fit", "fov_x_given", |
| "apply_mask", "force_projection", "call"))) |
| for r in m.warmed_variants)) |
| yield m |
| m.close() |
|
|
|
|
| def _real_images(n, size, seed): |
| """n distinct real images of ``size`` (w, h): random crops (and flips) of the demo photo.""" |
| src = Image.open(DEMO).convert("RGB") |
| rng = np.random.default_rng(seed) |
| W, H = src.size |
| out = [] |
| for i in range(n): |
| cw = int(W * rng.uniform(0.55, 0.95)) |
| ch = int(H * rng.uniform(0.55, 0.95)) |
| x0, y0 = int(rng.integers(0, W - cw + 1)), int(rng.integers(0, H - ch + 1)) |
| im = src.crop((x0, y0, x0 + cw, y0 + ch)).resize(size, Image.BICUBIC) |
| if i % 2: |
| im = im.transpose(Image.FLIP_LEFT_RIGHT) |
| out.append(np.asarray(im)) |
| return out |
|
|
|
|
| def _inputs(kind, size, seed, tmp_path): |
| imgs = _real_images(N_STEADY + 1, size, seed) |
| if kind == "array": |
| return imgs |
| if kind == "float": |
| return [torch.from_numpy(a).permute(2, 0, 1)[None].float().div_(255) for a in imgs] |
| paths = [] |
| for i, a in enumerate(imgs): |
| p = tmp_path / f"{kind}_{seed}_{i}.{'png' if kind == 'png' else 'jpg'}" |
| Image.fromarray(a).save(p, **({} if kind == "png" else {"quality": 92})) |
| paths.append(str(p)) |
| return paths |
|
|
|
|
| CANVAS = (1920, 1080) |
| VARIANTS = [ |
| ("array_canvas", "array", CANVAS, {}, "call"), |
| ("float_canvas", "float", CANVAS, {}, "call"), |
| ("png_canvas", "png", CANVAS, {}, "call"), |
| ("jpeg_800x600", "jpeg", (800, 600), {}, "call"), |
| ("stretch_800x600", "array", (800, 600), {"fit": "stretch"}, "call"), |
| ("fov_x", "array", CANVAS, {"fov_x": 60.0}, "call"), |
| ("no_mask_no_projection", "array", CANVAS, {"apply_mask": False, "force_projection": False}, "call"), |
| ("forward_float", "float", CANVAS, {}, "forward"), |
| ] |
|
|
|
|
| @pytest.mark.parametrize("name, kind, size, opts, method", VARIANTS, ids=[v[0] for v in VARIANTS]) |
| def test_first_call_is_fast(model, tmp_path, name, kind, size, opts, method): |
| xs = _inputs(kind, size, seed=zlib.crc32(name.encode()) % 1000, tmp_path=tmp_path) |
| fn = model.forward if method == "forward" else (lambda x: model(x, **opts)) |
| lat, outs = [], [] |
| for x in xs: |
| t = time.perf_counter() |
| outs.append(fn(x)) |
| lat.append(1e3 * (time.perf_counter() - t)) |
| first, steady = lat[0], statistics.median(lat[1:]) |
| print(f"\n[first-call] {name:24s} 1st {first:6.1f} ms 2nd {lat[1]:6.1f} ms 10th {lat[9]:6.1f} ms " |
| f"steady median {steady:6.1f} ms ratio {first / steady:.2f}") |
| again = fn(xs[0]) |
| if method == "forward": |
| for k in outs[0]: |
| assert torch.equal(outs[0][k], again[k]), k |
| else: |
| for k in outs[0].keys(): |
| np.testing.assert_array_equal(np.asarray(outs[0][k]), np.asarray(again[k])) |
| assert first <= steady * (1 + TOL) + SLACK_MS, (name, lat) |
|
|
