# 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.