litert-community/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.

Files

File	Size	Description
moge.tflite	136 MB	FP32 single-graph model, GPU-compatible

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.

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