| # MoGe-2 on Blackhole: Python API |
|
|
| `tt_moge.MoGeModel` runs MoGe-2 (ViT-L, normal head) on one Tenstorrent Blackhole chip from Python. |
| One RGB image goes in. Metric depth, a camera-space point map, surface normals, a validity mask and the camera intrinsics come out, at the input resolution. |
|
|
| ## Install |
|
|
| You need an environment that already has `ttnn` (tt-metal) and `torch`. |
|
|
| - In the tt-model container (`tt-model serve` image), `tt_moge` is already importable. Do not install anything. |
| - In a tt-metal Python environment, install the package. Run these commands from the repo root (the directory that contains `code/` and `media/`): |
|
|
| ```bash |
| pip install -e code/ # the Python API only (or: pip install code/) |
| pip install -e "code/[server]" # + the HTTP server: fastapi, uvicorn, pydantic, isal |
| pip install -e "code/[test]" # + pytest |
| pip install -e "code/[server,test]" # what the tests need: the host tests compare the API with the server |
| ``` |
|
|
| If your copy of `code/` has no `pyproject.toml`, `export PYTHONPATH=$PWD/code` works too. |
| The package declares only its own dependencies (`numpy<2`, `pillow`, `scipy`, `huggingface_hub`, `opencv-python-headless`). |
| The first use builds a small C++ host library with `g++` and caches it. |
| Without `g++`, the model uses the torch host code. The results are the same; the host part is slower. |
|
|
| ## Quickstart |
|
|
| ```python |
| from tt_moge import MoGeModel |
| |
| with MoGeModel.from_pretrained() as model: # opens device 0, gets the weights from the HF cache |
| out = model("photo.jpg") # a path, PIL image, numpy array or torch tensor |
| |
| print(out) # MoGeOutput(1920x1080, valid 90.5%, median depth 13.322 m, fov_x 86.1 deg, metric_scale 10.3329) |
| depth = out.depth # (H, W) float32, metres; inf where out.mask is False |
| points = out.points # (H, W, 3) float32, metres, camera space (x right, y down, z forward) |
| normal = out.normal # (H, W, 3) float32, unit normals |
| K = out.intrinsics_pixels # (3, 3) camera matrix in pixels |
| ``` |
|
|
| `python code/examples/quickstart.py` runs this snippet on the demo image `media/source.png`. |
| The script finds the demo image relative to its own file, so you can start it from any directory. |
| It writes `depth.png`, `normal.png`, `result.npz` and `result.json` to `./moge_output/` in the current directory (change it with `--out`). |
| Other relative paths in this document (`media/source.png`, `code/...`) start at the repo root. |
|
|
| `from_pretrained` loads the weights, compiles the device kernels, captures the device trace and warms up the host paths of the common calls (see [Warm-up](#warm-up)). |
| It took 42 s on a first run (kernel compile) and about 10 s with a warm kernel cache. |
| After that, the first call is as fast as the later calls: about 25 ms for a 1920x1080 image array (see [Speed](#speed)). |
| Keep the model object and call it many times. |
|
|
| ## API reference |
|
|
| ### `MoGeModel.from_pretrained(...)` |
| |
| ```python |
| MoGeModel.from_pretrained( |
| pretrained_model_name_or_path="Ruicheng/moge-2-vitl-normal", |
| *, device=None, device_id=0, revision=<pinned sha>, |
| resolution=(1920, 1080), num_tokens=1800, |
| dispatch="eth", num_command_queues=None, warmup=2, |
| warmup_variants="default", verbose=False) |
| ``` |
| |
| | Argument | Meaning | |
| |---|---| |
| | `pretrained_model_name_or_path` | A Hugging Face repo id, a directory that holds `model.pt`, or the path of `model.pt`. `$MOGE_WEIGHTS_DIR` overrides the repo. | |
| | `device` | An open ttnn device. The model uses it and does not close it. If you opened it with 2 command queues (worker dispatch only), also give `num_command_queues=2`. | |
| | `device_id` | The ttnn device to open when `device` is None. | |
| | `revision` | The weights revision on the Hub. The default is the revision of the published numbers. `None` = the default branch. | |
| | `resolution` | The canvas `(width, height)`. The model fits every image onto this canvas. Only `(1920, 1080)` is validated. | |
| | `num_tokens` | The ViT token budget, 1200-3600. Only 1800 is validated. | |
| | `dispatch` | `"eth"` (default): ETH dispatch, 1 command queue, 12x10 compute grid. This is the p150 configuration of the published numbers. ETH dispatch needs the ETH-dispatch patch of tt-metal. If the ETH open fails, the model gives a `RuntimeWarning` and uses worker dispatch. On a p150, worker dispatch gives an 11x10 compute grid, so the model is slower. `"worker"`: worker dispatch (opt-in). It gives 12x10 only on a Galaxy chip, 11x10 on a p150. | |
| | `num_command_queues` | 1 (default). 2 is an opt-in for worker dispatch only (Galaxy only): the output is read back on the second queue. The outputs are bit-identical with 1 or 2 queues. ETH dispatch supports 1 queue only. | |
| | `warmup` | Device forward passes at load time. The first one compiles the kernels and captures the trace. | |
| | `warmup_variants` | The call variants to prepare at load time, so that their first real call is fast. A preset (`"default"`, `"all"`, `"minimal"`, `"none"`), a variant dict, or a list of them. See [Warm-up](#warm-up). | |
| | `verbose` | `False` (default): ttnn is imported and the device is opened with `TT_LOGGER_LEVEL=error` and `LOGURU_LEVEL=WARNING`, so the tt-metal info and debug lines do not show. Errors still show. If you set these variables yourself, your values are used. `True`: progress lines (`[tt_moge] ...`, with the time of each warm-up variant) and the normal tt-metal log. | |
|
|
| The resolution and the token budget are fixed at load time, because the device trace is captured for them. |
| To use other values, load a second model. |
|
|
| ### `model(image, *, fit="pad", fov_x=None, apply_mask=True, force_projection=True)` |
| |
| Returns a `MoGeOutput` at the input resolution. A list of images returns a list of `MoGeOutput`. |
| |
| | Argument | Meaning | |
| |---|---| |
| | `image` | A file path; a PIL image; a numpy array `(H, W, 3)` uint8 or float in [0, 1] (also `(3, H, W)` and grayscale `(H, W)`); a torch tensor `(3, H, W)` or `(1, 3, H, W)` float in [0, 1], or `(H, W, 3)` uint8. Minimum size 8x8. | |
| | `fit` | `"pad"` puts the image on the canvas with gray bars and keeps the aspect ratio. `"stretch"` resizes the image to the canvas. | |
| | `fov_x` | The horizontal field of view of the image in degrees, if you know it. Then the model solves only for the depth shift. | |
| | `apply_mask` | Set invalid pixels to `inf` (depth, points) and 0 (normal). | |
| | `force_projection` | Compute the points from the depth and the intrinsics (upstream default). | |
| |
| Float images that do not have the canvas size are converted to 8 bit before the canvas fit, as in the HTTP server. |
| The model does not apply the EXIF orientation of JPEG files. |
| |
| ### `MoGeOutput` |
| |
| | Field | Type | Content | |
| |---|---|---| |
| | `points` | float32 `(H, W, 3)` | Point map in metres, camera space, OpenCV axes (x right, y down, z forward). | |
| | `depth` | float32 `(H, W)` | Metric depth (z) in metres. | |
| | `normal` | float32 `(H, W, 3)` | Unit surface normals in camera space. | |
| | `mask` | bool `(H, W)` | True for valid pixels. Sky and the letterbox bars are False. | |
| | `intrinsics` | float32 `(3, 3)` | Normalized intrinsics: `fx`, `cx` in units of the image width, `fy`, `cy` in units of the image height. | |
| | `intrinsics_pixels` | float32 `(3, 3)` | The same intrinsics in pixels. | |
| | `metric_scale` | float | The predicted scale. It is already applied to `points` and `depth`. | |
| | `fov_x`, `fov_y` | float | Field of view in degrees. | |
| |
| `out["depth"]`, `out.keys()` and `out.to_dict()` work like the upstream `infer()` dict. |
| `out.to_torch()` returns the upstream dict of torch tensors. |
| `out.save_npz(path)` writes the server's `npz` layout (`normal` float16, `mask` uint8). |
| |
| ### Other methods |
| |
| | Method | Use | |
| |---|---| |
| | `model.predict_many(images, **options)` | Several images, in order. The host loads the next image and post-processes the previous result while the device runs the current image. Same results as single calls. | |
| | `model.infer(image, fov_x=None, force_projection=True, apply_mask=True)` | Drop-in for the upstream `moge.model.v2.MoGeModel.infer`: a dict of torch tensors `points`, `intrinsics`, `depth`, `mask`, `normal`. A `(B, 3, H, W)` input gives batched outputs. `num_tokens` must be None or the load-time value; `resolution_level` and `use_fp16` are ignored. | |
| | `model.forward(image)` | The raw forward on a canvas-sized image, like the upstream `MoGeModel.forward`: `points` (affine, before the focal and shift recovery), `normal`, `mask` (probability), `metric_scale`, at the canvas resolution. This is the call that the published per-call latency measures. | |
| | `model.warmup(**variant)` | Prepares one more call variant (see [Warm-up](#warm-up)). Idempotent. Returns the seconds it took. | |
| | `model.warmed_variants` | The prepared variants and the warm-up time of each. | |
| | `model.close()` | Releases the trace and the device buffers. Closes the device if the model opened it. A `with` block calls it for you. | |
| |
| Calls are thread-safe. The device runs one image at a time (batch 1). |
| |
| ## Warm-up |
| |
| The device trace is captured for one canvas and one token budget. Every request uses the same trace, so there is only one device variant. |
| What changes per request is the host path: the decoder (array, float tensor, PNG file, JPEG file, PIL image), the input size (canvas fit and resize back), `fit`, `fov_x`, `apply_mask`, `force_projection`, and the method (`model(...)`, `forward`, `infer`, `predict_many`). |
| The first use of a host path pays one-time costs: lazy imports, decoder set-up, the first use of the C++ helpers, scratch buffers and allocator growth. |
| `from_pretrained` runs each warm-up variant once, with a synthetic photo-like image of the right size and kind (and a temporary PNG or JPEG file for the decoders). |
| Then the first real call of a warmed variant is as fast as the later calls. |
| No result is kept: the warm-up does not change the outputs. |
| |
| ### Options |
| |
| ```python |
| model = MoGeModel.from_pretrained(warmup_variants="default") # the default |
| model = MoGeModel.from_pretrained(warmup_variants="all") # also PIL, float 800x600, infer, predict_many, ... |
| model = MoGeModel.from_pretrained(warmup_variants=["default", {"input": "jpeg", "size": (4032, 3024)}]) |
| model = MoGeModel.from_pretrained(warmup_variants="none") # device warm-up only |
| model.warmup(input="png", size=(1280, 720), fit="stretch") # later; idempotent |
| ``` |
| |
| | Variant key | Values | Default | |
| |---|---|---| |
| | `input` | `"array"` (uint8 numpy), `"float"` (float32 torch `(1, 3, H, W)`), `"png"`, `"jpeg"` (a file path), `"pil"` | `"array"` | |
| | `size` | `(width, height)` of the input image | the canvas size | |
| | `fit`, `fov_x`, `apply_mask`, `force_projection` | as in `model(image, ...)`; for `fov_x` only "given or not" matters | the call defaults | |
| | `call` | `"call"`, `"predict_many"` (2 images), `"infer"`, `"forward"` (canvas size only) | `"call"` | |
| |
| | Preset | Variants | Host warm-up time | |
| |---|---|---:| |
| | `"default"` | uint8 array 1920x1080; float tensor 1920x1080; PNG file 1920x1080; JPEG file 800x600; uint8 800x600 with `fit="stretch"`; `fov_x` given; `apply_mask=False, force_projection=False`; `forward` on a float canvas | 0.74 s | |
| | `"all"` | `"default"` + PIL 800x600, float 800x600, `infer`, PNG 800x600, `predict_many` of PNG files, `apply_mask=False`, `force_projection=False`, `forward` on a uint8 canvas | 1.63 s | |
| | `"minimal"` | uint8 array at the canvas size | 0.27 s | |
| | `"none"` | no host warm-up (the 2 device forwards still run) | 0 s | |
| |
| Measured on chip 7 (warm kernel cache): `from_pretrained` took 9.5 s with `"none"` and 10.7 s with `"default"`; the A/B table below gives 9.7-10.6 s before and 9.8-12.4 s after on a loaded host. |
| A larger variant costs more: `model.warmup(input="png", size=(4032, 3024))` took 2.1 s (most of it is making and encoding the 12 MP synthetic image). |
| |
| A variant costs one full call (about 25-35 ms) plus the synthetic image (about 0.1 s for 1920x1080, made once per size per warm-up) and, for `png` / `jpeg`, the file encode. |
| The first variant also pays the host one-time costs (about 0.2 s of the 0.27 s of `"minimal"`). |
| The warm-up adds no device memory and no trace: all variants use the same trace. |
| `verbose=True` prints the time of each variant; `model.warmed_variants` lists them. |
| |
| ### What a warm-up cannot prepare |
| |
| - **Input size.** Any size from 8x8 works with the same trace, so the sizes are unbounded. A new size needs no compile. It only allocates its output arrays (about 60 MB of float32 maps at 1920x1080). Measured: with images of a different size in every call (`u8_native` below), the first call is 1.01x of the median. |
| - **PNG files larger than the largest warmed PNG.** The decoder grows its scratch buffer once, on the first larger image (one allocation of about 3 bytes per pixel; its time is not measured separately). Add a variant of that size if it matters (`model.warmup(input="png", size=(W, H))`). |
| - **Calls from a new thread.** The PNG decoder keeps its scratch buffer per thread. The first PNG call on a new thread allocates it again (`predict_many` decodes on its own short-lived thread, for every call). |
| - **The page cache.** A file that the OS has not read yet costs a disk read on every call, not only on the first one. |
| |
| ### Measured first call |
| |
| Fresh process for every row, `from_pretrained()` with the defaults, then 20 calls, each with a NEW real image (18 photos and their mirror images; never the warm-up image). |
| Median of 3 fresh processes per row, before (HEAD 2fc9c77: 2 device forwards + one host post-processing pass at load) and after (`warmup_variants="default"`). |
| Chip 23 of a shared Blackhole Galaxy host (load average about 20), 12x10 compute grid, worker dispatch, 2 command queues (2026-10-04). |
| Re-check in the p150 configuration (ETH dispatch, 1 command queue, 12x10; chip 1, 2026-10-05, `tt_moge/tests/test_api_first_call.py`, one fresh process): `from_pretrained` 9.9 s; 1st call / steady median: uint8 1920x1080 26.8 / 25.7 ms (1.04), float 1920x1080 26.4 / 25.1 ms (1.05), PNG 1920x1080 32.7 / 33.6 ms (0.97), JPEG 800x600 30.7 / 28.1 ms (1.10), 800x600 `fit="stretch"` 28.1 / 26.4 ms (1.07), `fov_x` 25.3 / 25.2 ms (1.00), `apply_mask=False, force_projection=False` 25.4 / 25.0 ms (1.01), `forward` 19.4 / 20.0 ms (0.97). All 8 variants pass (limit: 10 % + 1.5 ms). |
| The outputs of all 20 calls are bit-identical before and after, for every row. |
| |
| | Variant | `from_pretrained` s | 1st call ms | 2nd ms | 10th ms | steady median ms (calls 11-20) | 1st / steady | |
| |---|---:|---:|---:|---:|---:|---:| |
| | `model(uint8 array 1920x1080)` | 9.7 -> 10.0 | 26.2 -> 25.3 | 25.3 -> 25.4 | 25.9 -> 25.5 | 25.9 -> 25.1 | 1.01 -> 1.01 | |
| | `model(float tensor (1,3,1080,1920))` | 9.8 -> 10.1 | 26.0 -> 25.7 | 25.1 -> 25.5 | 26.1 -> 24.9 | 25.4 -> 25.3 | 1.04 -> 0.99 | |
| | `model("1920x1080.png")` | 9.8 -> 10.5 | 36.6 -> 34.8 | 34.5 -> 34.8 | 33.8 -> 34.9 | 34.8 -> 36.1 | 1.06 -> 0.95 | |
| | `model("800x600.jpg")` | 9.8 -> 10.7 | 34.9 -> 28.8 | 27.9 -> 31.1 | 29.8 -> 28.3 | 28.8 -> 30.5 | **1.21 -> 0.91** | |
| | `model("800x600.png")` | 10.6 -> 10.4 | 32.3 -> 27.5 | 28.3 -> 28.8 | 27.7 -> 30.9 | 28.7 -> 30.0 | **1.13 -> 0.99** | |
| | `model(uint8 array 800x600)` | 10.6 -> 10.2 | 26.2 -> 24.8 | 26.5 -> 27.3 | 26.7 -> 25.1 | 25.9 -> 25.5 | 1.01 -> 1.00 | |
| | `model(uint8 array, a new size every call)` | 9.7 -> 10.6 | 26.0 -> 25.6 | 26.6 -> 26.1 | 22.6 -> 24.4 | 25.0 -> 26.0 | 1.08 -> 1.01 | |
| | `model(PIL image 800x600)` | 9.5 -> 10.2 | 26.8 -> 26.1 | 25.9 -> 26.9 | 27.4 -> 26.1 | 28.5 -> 26.2 | 0.94 -> 1.00 | |
| | `model(float array 800x600)` | 10.1 -> 10.9 | 27.3 -> 26.8 | 27.1 -> 28.3 | 28.0 -> 27.2 | 28.6 -> 26.2 | 0.95 -> 1.00 | |
| | `model(800x600, fit="stretch")` | 10.1 -> 11.2 | 26.5 -> 28.8 | 26.1 -> 27.0 | 28.5 -> 28.8 | 26.8 -> 27.7 | 0.99 -> 1.01 | |
| | `model(1920x1080, fov_x=60)` | 11.4 -> 12.4 | 38.8 -> 27.5 | 32.1 -> 25.6 | 43.2 -> 26.2 | 47.8 -> 26.5 | 1.02 -> 1.08 | |
| | `model(1920x1080, apply_mask=False, force_projection=False)` | 10.6 -> 11.1 | 26.5 -> 24.7 | 26.8 -> 26.1 | 28.2 -> 27.7 | 26.3 -> 26.4 | 1.01 -> 0.97 | |
| | `model.forward(float canvas)` | 10.6 -> 10.1 | 19.3 -> 19.3 | 19.0 -> 19.5 | 20.2 -> 19.2 | 19.6 -> 19.2 | 0.98 -> 1.01 | |
| | `model.infer(float (3,600,800))` | 10.1 -> 9.8 | 26.7 -> 27.1 | 26.9 -> 27.9 | 27.3 -> 27.5 | 27.8 -> 26.6 | 0.95 -> 1.02 | |
| | `model.predict_many(4 x 800x600)`, per call of 4 | 9.6 -> 10.3 | 97.8 -> 96.3 | 98.7 -> 95.3 | 92.3 -> 95.6 | 94.0 -> 95.0 | 1.02 -> 1.01 | |
| |
| The load times and single calls vary by a few ms between processes on this shared host (other jobs run on it); the `fov_x` "before" processes ran during such a period. |
| Before the change, the first call of a file input was 6-21 % slower (decoder set-up); after it, every row is within 10 % of its steady median. |
| The steady latency did not change (`tt_moge/tests/test_api_device.py` after the change: `forward` median 20.03 ms, `model(array)` 26.2 ms, `model(path)` 34.4 ms; ETH, 1 CQ, chip 1, 2026-10-05: 20.28 / 25.56 / 34.19 ms). |
| The device test `tt_moge/tests/test_api_first_call.py` checks the first call of every default variant on a fresh model. |
| |
| ## Relation to the HTTP server |
| |
| `model(image)` does the same steps as `POST /predict` of `tt_moge/server/app.py`: canvas fit, device forward, focal and shift recovery, crop, resize to the input size, re-projection and mask. |
| For the same image and options, the arrays are bit-identical to the server's `npz` response (`tt_moge/tests/test_api.py`, `tt_moge/tests/test_api_device.py`). |
| The server stores `normal` as float16; the Python API gives float32. |
| |
| ## Speed |
| |
| Warm, batch 1, `media/source.png` (1920x1080), one Blackhole chip with a 12x10 compute grid, ETH dispatch, 1 command queue (the p150 configuration, the default since 2026-10-05). |
| |
| Measured on chip 1 of a shared Blackhole Galaxy host (2026-10-05, 2 processes of `perf/api/api_bench.py`, `perf/eth1cq/api_bench_eth1cq_*.log`), 30-60 warm calls per row. |
| The column "worker, 2 CQ" gives the earlier numbers (chip 16, 2026-10-04: worker dispatch, 2 command queues, a Galaxy-only configuration). |
|
|
| | Call | Median | Min | Median, worker, 2 CQ | |
| |---|---:|---:|---:| |
| | `model.forward(image)`, float32 canvas (the published `tt(image)` call) | 20.3-20.5 ms | 19.5-19.9 ms | 20.0-20.2 ms | |
| | `python -m tt_moge.bench_breakdown` (the published method; chip 1, 5 processes) | 20.3-21.1 ms | 19.2-19.5 ms | 20.1 ms | |
| | `model(array)`, uint8 1920x1080 (forward + focal and shift solve + finish at the input size) | 25.2-25.5 ms | 23.6-24.2 ms | 24.4-26.1 ms | |
| | `model(array)`, uint8 800x600 | 25.1-25.7 ms | 24.5-24.9 ms | 25.1-25.3 ms | |
| | `model("source.png")` (adds the PNG decode) | 33.8-34.3 ms | 31.9-32.9 ms | 33.6-36.1 ms | |
| | `model.predict_many(paths)`, per image (the next decode overlaps the device) | 26.1-26.3 ms | | 28.0-28.6 ms | |
| | `model.predict_many(arrays)`, per image | 25.5-25.8 ms | | 24.4-25.7 ms | |
|
|
| On chip 1, worker dispatch + 2 command queues gives the same `bench_breakdown` e2e as ETH + 1 queue (20.5-20.9 vs 20.3-21.1 ms, alternated): the difference to chip 16 / chip 9 is the chip and the host load, not the configuration. |
| The published `tt(image)` figure (README, OPT_REPORT.md "p150 ETH-dispatch compliance 2026-10-05", chip 1) is 20.3-21.1 ms median, 19.2-19.5 ms min. |
| `model(image)` adds the host post-processing that `tt(image)` does not include: about 2 ms for the focal and shift solve and about 3 ms for the finish pass (16 host threads, `MOGE_SERVE_THREADS`). |
| The server's `timing_ms.device` stage is the forward part only. |
|
|
| ## Tests |
|
|
| Run the tests from `code/`. The device tests find the demo image relative to their own file. |
|
|
| ```bash |
| cd code |
| pip install -e ".[server,test]" # once: pytest + the server dependencies the host tests compare against |
| |
| # host only: no device, no ttnn device open, no weights (about 1 min) |
| TT_VISIBLE_DEVICES=none python -m pytest tt_moge/tests/test_api.py tt_moge/tests/test_serve_fastpath.py -q |
| |
| # one chip each; run each file in its own process (each opens the device) |
| python -m pytest tt_moge/tests/test_api_device.py -q -s # API == server npz, bit for bit; call latency |
| python -m pytest tt_moge/tests/test_api_first_call.py -q -s # first real call after from_pretrained is fast |
| ``` |
|
|
| Without the `server` extra, the API-vs-server comparisons in `test_api.py` are skipped. |
| Set `MOGE_DEVICE` to choose the device id (default 0). `MOGE_FIRST_CALL_TOL` (default 0.10) sets the allowed first-call ratio. |
|
|
