code/PYTHON.md

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/):

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

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). 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). Keep the model object and call it many times.

API reference

MoGeModel.from_pretrained(...)

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

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.

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.