pablovela5620/moge-v2-tensorrt-zerogpu

MoGe v2: PyTorch versus TensorRT on ZeroGPU

This is a deliberately small reference Space. Ten video frames go through the same fixed MoGe v2 ViT-L normal/mask architecture and weights in PyTorch and TensorRT. Each frame streams into an embedded Rerun viewer on a frame timeline. Gradio reports synchronized latency, setup time, and numerical agreement beside it.

Files

• —app.py: the standard Hugging Face entrypoint; it only installs and enters the locked Pixi environment.
• —moge_app.py: the ZeroGPU callback, Gradio/Rerun UI, preprocessing, timing, PyTorch, and the fixed-shape TensorRT runner.
• —pixi.toml and pixi.lock: the complete CUDA 13, PyTorch 2.11, TensorRT 11.2.1.2, Rerun 0.35, Spaces, and Gradio environment.
• —moge.engine: the prebuilt TensorRT plan.
• —example.mp4: ten sample frames derived from an image in the official MoGe repository.

The source-free, explicit-FP16 PyTorch ExportedProgram baseline is pinned in `pablovela5620/moge-v2-normal-trt-assets`. It is fetched during CPU startup, outside the ZeroGPU allocation. Keeping the 642 MiB PyTorch artifact separate leaves room for the 636 MiB engine under the Space repository's 1 GB limit.

There is no ONNX graph, compiler, model source tree, or engine build path in the Space. Building happens once, outside ZeroGPU; serving only loads the pinned artifacts.

Gradio and Rerun compatibility

Keep gradio==6.20.0, gradio-rerun==0.35.0, and rerun-sdk==0.35.0 together. The official gradio-rerun 0.35 lock also resolves Gradio 6.20.0.

Gradio 6.21.0 introduced a frontend regression that is still present in 6.26.0. Its app-tree update change mirrors each output update into reactive node props before calling the custom component's normal update callback. That retriggers Gradio's custom-component mount effect: the Rerun canvas is removed, its cleanup stops the WebViewer, and queued RRD chunks then target the stopped viewer. The browser reports attempted to open channel "gradio" in a stopped web viewer and repeated send_rrd errors. A minimal local bisection found 6.20.0 good and 6.21.0 through 6.26.0 bad.

Four details are also deliberate:

• —@spaces.GPU is the outer decorator, so ZeroGPU owns the whole streamed job.
• —@rr.thread_local_stream(...) creates an isolated recording for each generator invocation. The callback uses global rr.* calls and yields rr.binary_stream().read() after every frame.
• —Rerun(streaming=True) receives those incremental RRD bytes, matching the official component's streaming contract.
• —launch(ssr_mode=False) matches the official examples and avoids Hugging Face Spaces' default Gradio SSR path for this WebAssembly viewer.

Reproducibility contract

• —Model: Ruicheng/moge-2-vitl-normal at revision b135031bae30b5ac2ae141a0e68717795ce38340
• —Outputs: camera-space normal and mask heads only
• —Input: float32 RGB [1, 3, 756, 1008] in [0, 1]
• —Tokens: 3600
• —PyTorch: 2.11.0, CUDA 13.0, explicit FP16 weights/compute with float32 I/O
• —PyTorch artifact SHA-256: 908ab397abf18938d681b503a8be4cbfd1cced6fd1fa6fb67855d5cea0e851f5
• —TensorRT: 11.2.1.2, strongly typed, SM 12.0 same-compute-capability compatibility
• —Engine profile: batch min 1, optimum 2, maximum 2
• —Engine SHA-256: bd109885ad26fc2c77d7feb54a68c1c2b31b8afbab65e1162c851ac516f8ffc1

The engine was built on an RTX 5090 in 28.71 seconds. Its SM 12.0 compatibility mode also loaded on ZeroGPU's default RTX PRO 6000 Blackwell. A different compute capability or TensorRT version needs a new engine.

Verified results

These are synchronized batch-1 measurements of the deployed artifacts at 756 × 1008. Each backend gets two warmups, followed by one timed inference per video frame. Setup and rendering are excluded.

On the local ten-frame run, the mean angular difference was 0.531° and mask mean absolute error was 1.92e-10. PyTorch setup took 2.76 seconds and TensorRT engine deserialization took 0.64 seconds. The UI reports setup separately so it is not mistaken for steady-state inference.

Run

On a CUDA 13 host:

CONDA_OVERRIDE_CUDA=13.0 pixi run --frozen app

Hugging Face starts app.py in its standard image. That launcher pins and installs Pixi 0.73.0 if needed, then runs the Pixi app task. The task launches moge_app.py inside the frozen environment, so application dependencies still come only from the committed Pixi lock.

The bare @spaces.GPU decorator lets ZeroGPU choose its default allocation. FFmpeg decoding runs in the preceding CPU callback. Model and engine setup happen once inside the streamed GPU job, and the UI reports them separately. Each frame gets one synchronized CUDA-event measurement per backend.

See the ZeroGPU documentation and GPU size selection.