MoGe-2-LiteRT

MoGe-2 ViT-S — LiteRT (TFLite) GPU

On-device LiteRT (.tflite) conversion of MoGe-2 (CVPR'25 Oral) monocular geometry estimation, converted from Ruicheng/moge-2-vits-normal (DINOv2 ViT-S backbone, 35M params).

A single forward pass turns one RGB image into an affine 3D point map, surface normals, a confidence mask, and a metric scale — enabling depth, surface normals, and a rotatable 3D point cloud on a phone.

The model runs fully on the LiteRT CompiledModel GPU accelerator (ML Drift): all 836 ops are GPU-native, no CPU fallback, no Flex ops.

Try it in your browser: google-ai-edge.github.io/litert-samples/…/moge — this model running on LiteRT.js (WebGPU runs moge_fp16.tflite, the WASM fallback moge.tflite). Nothing to install; inference runs on your machine.

Files

moge_fp16.tflite halves the download for GPU use and matches FP32 up to the weight cast: across 7 test photos, worst-case max |Δ| is 0.0027 on points and 0.0026 on normal, mask decisions at 0.5 agree on ≥ 99.997% of pixels, and scale is within 0.05%. On WebGPU (LiteRT.js, M4 Mac) it also runs slightly faster (55.6 vs 58.5 ms/photo). Keep FP32 for CPU/WASM: XNNPACK declines the FP16 graph and falls back to reference kernels (~31× slower).

I/O

• Input: [1, 3, 448, 448] float32, NCHW, RGB normalized to [0, 1] (ImageNet mean/std is applied inside the graph).
• Outputs (4):
  • points [1, 448, 448, 3] — affine point map (exp remap: [xy·exp(z), exp(z)])
  • normal [1, 448, 448, 3] — L2-normalized surface normals
  • mask [1, 448, 448, 1] — sigmoid confidence (> 0.5 = valid)
  • scale [1, 1, 1, 1] — metric scale factor

Usage (Android, LiteRT CompiledModel)

val model = CompiledModel.create(
    context.assets, "moge.tflite",
    CompiledModel.Options(Accelerator.GPU), null
)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()
inputs[0].writeFloat(nchwFloatArray)   // [1,3,448,448], RGB [0,1]
model.run(inputs, outputs)
val points = outputs[0].readFloat()    // identify the 4 outputs by element count + range

Python (desktop verification)

import numpy as np
from PIL import Image
from ai_edge_litert.interpreter import Interpreter

img = Image.open("photo.jpg").convert("RGB").resize((448, 448))
x = (np.asarray(img, np.float32) / 255.0).transpose(2, 0, 1)[None]
it = Interpreter(model_path="moge.tflite"); it.allocate_tensors()
it.set_tensor(it.get_input_details()[0]["index"], x); it.invoke()
outs = [it.get_tensor(o["index"]) for o in it.get_output_details()]
# identify outputs by shape/range: `normal` is the [1,448,448,3] tensor
# whose vectors have unit L2 norm; `mask` > 0.5 marks valid pixels.

A complete Android sample (gallery → normal map + depth) is available in google-ai-edge/litert-samples.

Performance

• ~522 ms / frame on a Pixel 8a (Mali-G615) GPU.

Conversion notes

Converted with litert-torch (NCHW preserved — required for ViT attention accuracy). Making DINOv2 + the ConvStack decoder fully GPU-compatible required nine graph rewrites (LayerScale bake, fused-qkv decomposition, position-embedding bake, ConvTranspose → bilinear+1×1, etc.). Verified: all ops GPU-native, output correlation ≈ 1.0 vs. the PyTorch reference.

Performance

Measured on a Pixel 8a (Tensor G3, Android 16) with the standard TFLite benchmark_model tool — 10 warm-up runs then 50 timed runs, reported as the tool's mean.

The two GPU rows are different runtimes, not a contradiction. The LITERT_CL figure is the one recorded when this model shipped, taken through LiteRT's own CompiledModel accelerator — the path the Kotlin sample app and the LiteRT API use. The TfLiteGpuDelegateV2 figure is the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. They agree on how much of the graph the GPU takes; they disagree on speed, and the classic delegate is the slower of the two here. Read the TfLiteGpuDelegateV2 row as a reproducible floor, not as this model's speed on LiteRT.

Note that the GPU does not take the whole graph here (169 / 835); the remainder runs on the CPU and the split costs a per-partition round trip.

Snapdragon NPU (Hexagon)

The NPU is 1.90x faster than the GPU (56.22 ms against 106.8 ms) and loads 9.13x faster (227 ms against 2071 ms).

Measured on a Samsung Galaxy S26 (Snapdragon 8 Elite Gen 5 / SM8850, Hexagon v81, Android 16), LiteRT CompiledModel 2.2.0, one accelerator per process, 5 warm-up runs then N=50 timed runs, median reported. Every run held thermal status NONE throughout. Headroom 0.71-0.75, where 1.0 is the throttling threshold.

The NPU rows here ran artifacts compiled ahead of time for SM8850 with QAIRT 2.47.0; the GPU rows ran the published files as they are. LiteRT can also compile for the NPU on the device at first load, which is what lets you ship the published file unchanged — that path and the ten runtime libraries it needs are in the NPU recipe, and we did not measure it here. GPU wiring is in the GPU recipe.

Raspberry Pi 5 (CPU)

Measured on a Raspberry Pi 5 Model B Rev 1.1 (8 GB, Raspberry Pi OS 64-bit) with the LiteRT benchmark_model tool from litert-cli-nightly 0.2.0.dev20260805: CPU inference (XNNPACK, 4 threads), 3 invocations per file of 10 warm-up plus 50 timed runs (the tool caps a phase at 150 s, so very slow graphs run fewer — the Runs column is the actual timed total). The latency is the median across invocations; the spread is the min–max over all timed runs. No thermal throttling occurred during these runs (vcgencmd get_throttled stayed 0x0).

License & attribution

• Model: MIT (original microsoft/MoGe).
• DINOv2 backbone components: Apache-2.0.
• This is a format conversion of Ruicheng/moge-2-vits-normal; all credit to the original authors (Microsoft Research).