SAM 3D Body on TensorRT: an end-to-end INT8 deployment tutorial

SAM 3D Body is Meta’s single-image full-body human-mesh-recovery model (CVPR 2026). It pairs a DINOv3 ViT-H+/16 encoder (840 M parameters — the dominant cost) with a promptable decoder + MHR parametric head that regresses an 18439-vertex human mesh.

This tutorial deploys the SAM 3D Body model with TensorRT for fast inference on NVIDIA GPUs using embedl-deploy-tensorrt. The pipeline consists of:

1. Fetching the HuggingFace model and GitHub code.

2. Quantizing the encoder with the Embedl-Deploy PTQ pipeline.

3. Making the decoder + MHR head torch.onnx-exportable.

4. Building strongly-typed TensorRT engines for both with TensorRT 11.

5. Running a real image end-to-end on the engines, verifying fidelity against eager PyTorch, benchmarking, and rendering the recovered 3D mesh.

Results on an NVIDIA L4 (TensorRT 11.1, 512x512 input, batch 1):

encoder  INT8+FP16   ~34 ms      (FP32 eager PyTorch: ~237 ms -> 7.0x)
decoder  FP16        ~16 ms
total                ~50 ms

full-pipeline mesh error vs FP32 eager : 2.3 mm mean / 9.7 mm max

Requirements (an NVIDIA GPU with ~20 GB VRAM; ~15 GB of disk for artifacts):

pip install torch tensorrt pycuda numpy pillow opencv-python matplotlib
pip install imageio onnx huggingface_hub
pip install pytorch-lightning yacs scikit-image einops
pip install timm roma loguru optree termcolor
pip install braceexpand omegaconf onnxscript
pip install "embedl-deploy["tensorrt"]"

The facebook/sam-3d-body-dinov3 checkpoint is HuggingFace-gated: accept the license on the model page and log in (hf auth login) first.

This tutorial ships as two files that must sit in the same directory: this script and its export-utilities companion sam_3d_body_export.py. Note that the download buttons at the bottom of this page bundle only this script — fetch the companion module through the link above. Download both, then run:

python sam_3d_body.py

Every stage caches its artifact in ./artifacts/ and is skipped when the artifact already exists, so an interrupted run resumes where it stopped.

Stage 0 — configuration and public sources

Everything the pipeline needs is public:

• facebook/sam-3d-body-dinov3 (gated): model.ckpt with all weights, and assets/mhr_model.pt, the TorchScript MHR body rig.

• facebookresearch/dinov3 (GitHub): the encoder architecture code, loaded through torch.hub with pretrained=False — the encoder weights come from the SAM-3D-Body checkpoint, not the DINOv3 release.

• facebookresearch/sam-3d-body (GitHub): the decoder/head code.

import json
import statistics
import subprocess
import sys
import time
import types
import warnings
from collections.abc import Callable
from pathlib import Path

import cv2  # type: ignore[import-not-found]
import imageio.v2 as imageio  # type: ignore[import-not-found]
import matplotlib as mpl  # type: ignore[import-not-found]
import numpy as np  # type: ignore[import-not-found]
import numpy.typing as npt  # type: ignore[import-not-found]
import onnx
import pycuda.autoinit  # type: ignore[import-not-found]  # noqa: F401
import pycuda.driver as cuda  # type: ignore[import-not-found]
import tensorrt as trt
import torch
from huggingface_hub import hf_hub_download  # type: ignore[import-not-found]
from onnx import TensorProto, numpy_helper
from PIL import Image  # type: ignore[import-not-found]
from torch import nn

from embedl_deploy import transform
from embedl_deploy.backend import set_backend

set_backend("tensorrt")
from embedl_deploy.quantize import (
    CalibrationMethod,
    ModulesToSkip,
    QuantConfig,
    TensorQuantConfig,
    quantize,
)
from embedl_deploy.tensorrt import TENSORRT_PATTERNS

mpl.use("Agg")
import matplotlib.pyplot as plt  # type: ignore[import-not-found]
from matplotlib.axes import Axes  # type: ignore[import-not-found]
from matplotlib.collections import (  # type: ignore[import-not-found]
    PolyCollection,
)

# the SAM-3D-Body code (decoder, MHR head, estimator utilities), pinned to
# a known commit so runs are reproducible and the imported code is exactly
# the reviewed revision
HERE = Path(__file__).resolve().parent
REPO_URL = "https://github.com/facebookresearch/sam-3d-body"
REPO_COMMIT = "b5c765a0d89d789985e186d396315e7590887b94"
DINOV3_COMMIT = "346f38fee679c56a6888f91c51670fae61d364e0"
REPO_DIR = HERE / "sam-3d-body"
if not REPO_DIR.exists():
    print(f"fetching facebookresearch/sam-3d-body @ {REPO_COMMIT[:12]} ...")
    subprocess.run(["git", "init", "-q", str(REPO_DIR)], check=True)
    subprocess.run(
        [
            "git",
            "-C",
            str(REPO_DIR),
            "fetch",
            "-q",
            "--depth",
            "1",
            REPO_URL,
            REPO_COMMIT,
        ],
        check=True,
    )
    subprocess.run(
        ["git", "-C", str(REPO_DIR), "checkout", "-q", REPO_COMMIT],
        check=True,
    )
repo_head = subprocess.run(
    ["git", "-C", str(REPO_DIR), "rev-parse", "HEAD"],
    check=True,
    capture_output=True,
    text=True,
).stdout.strip()
if repo_head != REPO_COMMIT:
    print(
        f"    warning: sam-3d-body checkout is at {repo_head[:12]}, "
        f"expected {REPO_COMMIT[:12]}"
    )
sys.path.insert(0, str(REPO_DIR))

from sam_3d_body import load_sam_3d_body  # type: ignore[attr-defined]
from sam_3d_body.models.heads import (  # type: ignore[import-not-found]
    mhr_head as MH,
)
from sam_3d_body.models.modules import (  # type: ignore[import-not-found]
    mhr_utils as MU,
)
from sam_3d_body_export import (
    _camera_encoder_fwd_nodownsample,
    _patched_embed_keypoints,
    _patched_mhr_head_forward,
    _patched_perspective_projection,
    _patched_replace_hands_in_pose,
    _safe_compact_cont_to_model_params_body,
    _safe_compact_cont_to_model_params_hand,
    flatten_rig_for_export,
    register_ts_exporter_symbolics,
)

warnings.filterwarnings("ignore")
torch.manual_seed(0)

ART = HERE / "artifacts"
ART.mkdir(exist_ok=True)

IMG_SIZE = 512  # SAM-3D-Body production resolution (32x32 tokens)
N_CALIB = 16  # PTQ calibration batches
SAMPLE_IMAGE_URL = "https://github.com/facebookresearch/sam-3d-body/raw/main/assets/qualitative_comparisons/sample1/input_bbox.png"
SAMPLE_IMG = HERE / "sample_input.png"

FP32_PT2 = ART / "sam3dbody_backbone_fp32.pt2"
REF_FEATURES = ART / "ref_features_fp32.npy"
ENC_ONNX = ART / "embedl_sam3dbody_int8.onnx"
DEC_ONNX = ART / "sam3dbody_decoder.onnx"
FACES_NPY = ART / "sam3dbody_faces.npy"
REF_MESH = ART / "ref_mesh_fp32.npz"
STATS_PATH = ART / "sam3dbody_trt_stats.json"

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], np.float32).reshape(3, 1, 1)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], np.float32).reshape(3, 1, 1)

DEC_INPUT_NAMES = [
    "features",
    "rays_ds",
    "condition_info",
    "keypoints",
    "bbox_center",
    "bbox_scale",
    "ori_img_size",
    "cam_int",
    "affine_trans",
    "img_size",
]
DEC_OUTPUT_NAMES = [
    "pred_vertices",
    "pred_cam_t",
    "focal_length",
    "pred_keypoints_3d",
    "pred_keypoints_2d",
    "global_rot",
    "body_pose",
    "scale",
    "shape",
    "hand",
    "pred_pose_raw",
    "pred_cam",
]

STATS: dict[str, float | int | str] = {}
print("=" * 70)
print(f"torch  : {torch.__version__}")
print(f"GPU    : {torch.cuda.get_device_name(0)}")

# the gated checkpoint (weights + TorchScript MHR rig)
print("fetching facebook/sam-3d-body-dinov3 (cached after first run) ...")
CKPT_PATH = hf_hub_download("facebook/sam-3d-body-dinov3", "model.ckpt")
MHR_PATH = hf_hub_download(
    "facebook/sam-3d-body-dinov3", "assets/mhr_model.pt"
)
# read by load_sam_3d_body
hf_hub_download("facebook/sam-3d-body-dinov3", "model_config.yaml")

if not SAMPLE_IMG.exists():
    print(f"fetching sample image {SAMPLE_IMAGE_URL} ...")
    imageio.imwrite(SAMPLE_IMG, imageio.imread(SAMPLE_IMAGE_URL))


def benchmark_ms(
    fn: Callable[[], object], n: int = 50, warmup: int = 10
) -> tuple[float, float]:
    """Mean and p95 wall-clock latency of `fn` in milliseconds.

    :param fn:
        The callable to time; it must block until its work completes
        (synchronize the GPU inside).
    :param n:
        Number of timed iterations.
    :param warmup:
        Number of untimed warmup iterations.
    :returns:
        The ``(mean, p95)`` latency in milliseconds.
    """
    for _ in range(warmup):
        fn()
    times = []
    for _ in range(n):
        t0 = time.perf_counter()
        fn()
        times.append((time.perf_counter() - t0) * 1e3)
    return statistics.mean(times), sorted(times)[int(0.95 * n)]

Preprocessing — a numpy port of the upstream top-down transforms

SAM 3D Body is a top-down method: it takes a person bounding box, expands it by 1.25x, fixes the aspect ratio, and warps the crop to 512x512. The decoder is additionally conditioned on the camera: a ray direction per feature cell (computed from a default full-image intrinsic matrix) and a CLIFF-style bbox-location vector. We reproduce those transforms in plain numpy so the deployed pipeline has no PyTorch dependency at inference time.

def read_rgb(path: Path) -> npt.NDArray[np.uint8]:
    """Load an image file as an RGB uint8 array."""
    img = cv2.imread(str(path))
    assert img is not None, f"cannot read image {path}"
    return np.asarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), dtype=np.uint8)


def _fix_aspect_ratio(
    scale: npt.NDArray[np.float32], ratio: float
) -> npt.NDArray[np.float32]:
    """Grow the smaller side of `scale` to match the aspect `ratio`."""
    w, h = scale
    return np.array(
        [w, w / ratio] if w > h * ratio else [h * ratio, h], np.float32
    )


def _third(
    a: npt.NDArray[np.float32], b: npt.NDArray[np.float32]
) -> npt.NDArray[np.float32]:
    """Return ``b + (a - b)`` rotated 90 degrees anticlockwise."""
    d = a - b
    return b + np.array([-d[1], d[0]], np.float32)


def _warp_matrix(
    center: npt.NDArray[np.float32],
    scale: npt.NDArray[np.float32],
    out: int = IMG_SIZE,
) -> npt.NDArray[np.float64]:
    """Upstream ``get_warp_matrix`` with rot=0 (bbox_utils)."""
    # numpy broadcast add, not list concatenation
    src = np.stack([center, center + [0, -0.5 * scale[0]]])  # noqa: RUF005
    src = src.astype(np.float32)
    dst = np.array([[out / 2, out / 2], [out / 2, 0.0]], np.float32)
    src = np.vstack([src, _third(src[0], src[1])])
    dst = np.vstack([dst, _third(dst[0], dst[1])])
    # float64, as upstream
    return np.asarray(cv2.getAffineTransform(src, dst), dtype=np.float64)


def preprocess(
    image_path: Path, bbox: tuple[int, int, int, int] | None = None
) -> tuple[
    npt.NDArray[np.float32],
    dict[str, npt.NDArray[np.float32]],
    npt.NDArray[np.uint8],
]:
    """Turn an image and person bbox into encoder and decoder inputs.

    :param image_path:
        Path to the input image.
    :param bbox:
        Person bounding box ``(x1, y1, x2, y2)``; the full image when
        omitted.
    :returns:
        The normalized ``(1, 3, 512, 512)`` encoder input, the decoder
        conditioning feed dict (everything except ``features``), and the
        RGB crop for visualization.
    """
    img = read_rgb(image_path)
    h, w = img.shape[:2]
    x1, y1, x2, y2 = bbox or (0, 0, w, h)

    center = np.array([(x1 + x2) / 2, (y1 + y2) / 2], np.float32)
    # GetBBoxCenterScale then TopdownAffine, as in the upstream pipeline
    scale = np.array([x2 - x1, y2 - y1], np.float32) * 1.25
    scale = _fix_aspect_ratio(_fix_aspect_ratio(scale, 0.75), 1.0)
    A = _warp_matrix(center, scale)
    crop = np.asarray(
        cv2.warpAffine(img, A, (IMG_SIZE, IMG_SIZE), flags=cv2.INTER_LINEAR),
        dtype=np.uint8,
    )
    x = (
        crop.transpose(2, 0, 1).astype(np.float32) / 255.0 - IMAGENET_MEAN
    ) / IMAGENET_STD

    # default full-image intrinsics: focal = image diagonal, principal point
    # at the image center (upstream's no-calibration fallback)
    focal = float(np.hypot(w, h))
    cam = np.array(
        [[[focal, 0, w / 2], [0, focal, h / 2], [0, 0, 1]]], np.float32
    )

    # camera-ray conditioning: the ray (x,y) for every crop pixel, antialias-
    # downsampled to the 32x32 feature grid (matches the upstream
    # CameraEncoder's F.interpolate(..., antialias=True))
    gx, gy = np.meshgrid(
        np.arange(IMG_SIZE), np.arange(IMG_SIZE), indexing="xy"
    )
    grid = np.stack([gx, gy], axis=2).astype(np.float32)
    A32 = A.astype(np.float32)
    diag, trans = A32[[0, 1], [0, 1]], A32[[0, 1], [2, 2]]
    grid = grid / diag - trans / diag  # crop pixel -> image pixel
    grid = (grid - cam[0, [0, 1], [2, 2]]) / cam[0, [0, 1], [0, 1]]  # -> ray
    n = IMG_SIZE // 16
    rays_ds = np.stack(
        [
            np.asarray(
                Image.fromarray(grid[:, :, c]).resize(
                    (n, n), Image.Resampling.BILINEAR
                )
            )
            for c in range(2)
        ]
    )[None]

    # CLIFF condition, intrinsic-center variant
    cond = np.array(
        [
            [
                (center[0] - cam[0, 0, 2]) / focal,
                (center[1] - cam[0, 1, 2]) / focal,
                scale[0] / focal,
            ]
        ],
        np.float32,
    )
    feeds = {
        "rays_ds": rays_ds.astype(np.float32),
        "condition_info": cond,
        "keypoints": np.array([[[0.0, 0.0, -2.0]]], np.float32),  # dummy
        "bbox_center": center.reshape(1, 1, 2),
        "bbox_scale": scale.reshape(1, 1, 2),
        "ori_img_size": np.array([[[w, h]]], np.float32),
        "cam_int": cam,
        "affine_trans": A32.reshape(1, 1, 2, 3),
        "img_size": np.full((1, 1, 2), IMG_SIZE, np.float32),
    }
    return x[None], feeds, crop


print("\n[stage 0] preprocessing the sample image ...")
enc_input, dec_feeds, crop_rgb = preprocess(SAMPLE_IMG)
print(
    f"    encoder input {enc_input.shape}, decoder conditioning "
    f"{sorted(dec_feeds)}"
)

Stage 1 — the encoder, with the real checkpoint weights

The checkpoint stores the whole model; the encoder lives under the backbone.encoder. prefix. We wrap it: plain tensors in, plain tensors out, which is what torch.export and the Embedl PTQ flow expect.

Why deploy the encoder separately from the decoder? The full SAM3DBody.forward mixes data-dependent control flow (body/hand batch selection), so it doesn’t fit a single static graph. The encoder is a clean image-in/features-out function and carries most of the compute, so it is the natural quantization target; the decoder is exported as a second, FP16-only graph in stage 3.

class BackboneWrapper(nn.Module):
    """Tensor-in/tensor-out wrapper around the DINOv3 encoder."""

    def __init__(self, encoder: nn.Module) -> None:
        """Wrap the DINOv3 `encoder`."""
        super().__init__()
        self.encoder = encoder

    def forward(self, image: torch.Tensor) -> torch.Tensor:
        """Return the final 32x32 feature map for `image`."""
        feats = self.encoder.get_intermediate_layers(  # type: ignore[operator]
            image, n=1, reshape=True, norm=True
        )
        return feats[-1]  # type: ignore[no-any-return]  # (B, 1280, 32, 32)


def load_encoder_with_real_weights(device: str) -> nn.Module:
    """Build the DINOv3 architecture and load the SAM-3D-Body weights.

    :param device:
        Device to place the encoder on.
    :returns:
        The encoder in eval mode with the checkpoint weights loaded.
    """
    encoder = torch.hub.load(  # type: ignore[no-untyped-call]
        f"facebookresearch/dinov3:{DINOV3_COMMIT}",
        "dinov3_vith16plus",
        source="github",
        pretrained=False,
        trust_repo=True,
        skip_validation=True,  # ref is a pinned commit, not a branch
    ).eval()
    ckpt = torch.load(CKPT_PATH, map_location="cpu", weights_only=False)
    state = ckpt.get("state_dict", ckpt)
    backbone_state = {
        k.removeprefix("backbone.encoder."): v
        for k, v in state.items()
        if k.startswith("backbone.encoder.")
    }
    missing, unexpected = encoder.load_state_dict(backbone_state, strict=False)
    # the checkpoint omits only the mask_token, which inference never uses
    assert set(missing) <= {"mask_token"}, f"missing encoder keys: {missing}"
    assert not unexpected, f"checkpoint keys not in encoder: {unexpected}"
    return encoder.eval().to(device)  # type: ignore[no-any-return]


if not (FP32_PT2.exists() and REF_FEATURES.exists()):
    print("\n[stage 1] loading the DINOv3 ViT-H+/16 encoder (real weights)...")
    model = BackboneWrapper(load_encoder_with_real_weights("cuda"))
    n_params = sum(p.numel() for p in model.parameters())
    print(
        f"    {n_params / 1e6:.1f}M params, {n_params * 4 / 1e9:.2f} GB FP32"
    )

    x = torch.from_numpy(enc_input).cuda()
    with torch.no_grad():
        ref = model(x)
    np.save(REF_FEATURES, ref.float().cpu().numpy())
    print(f"    FP32 reference features {tuple(ref.shape)} saved")

    # eager FP32 latency -- the baseline the TRT engine is measured against
    def _run_eager() -> None:
        with torch.no_grad():
            model(x)
        torch.cuda.synchronize()

    STATS["encoder_fp32_eager_ms"], _ = benchmark_ms(
        _run_eager, n=20, warmup=5
    )
    print(f"    eager FP32 latency: {STATS['encoder_fp32_eager_ms']:.1f} ms")

    print("    torch.export.export(encoder) ...")
    dummy = torch.randn(1, 3, IMG_SIZE, IMG_SIZE, device="cuda")
    with torch.no_grad():
        exported = torch.export.export(model, (dummy,), strict=False)
    torch.export.save(exported, str(FP32_PT2))
    print(f"    saved {FP32_PT2.name}")
    del model, exported, ref
    torch.cuda.empty_cache()
else:
    print("\n[stage 1] reusing cached FP32 export + reference features")

Stage 2 — INT8 post-training quantization with Embedl-Deploy

The standard Embedl-Deploy recipe — transform fuses the exported graph into TensorRT-friendly patterns, quantize inserts Q/DQ nodes and calibrates them — plus three customizations that matter for this transformer:

• Skip the patch-embed conv. The 3-channel input projection gains almost nothing from INT8 and loses accuracy disproportionately.

• Don’t SmoothQuant the LayerNorms. Smoothing folds scales into LayerNorm weights in-place; on this backbone that hurts more than it helps, so LayerNorm modules are excluded via ModulesToSkip(smooth=...).

• Calibrate on the real crop. MINMAX with 16 lightly-jittered copies of the actual person crop gives ranges that match deployment.

Note

To keep the tutorial self-contained, PTQ calibrates on jittered copies of a single image. In a realistic deployment you would calibrate on a few hundred representative person crops from your target domain — for example a sample of the datasets SAM 3D Body itself is trained and evaluated on, such as COCO, MPII, AI Challenger, 3DPW, EgoHumans, EgoExo4D, Harmony4D, and SA-1B (see data/scripts in the upstream sam-3d-body repository for download helpers).

Note that the final INT8 coverage is decided by the backend’s precision propagation inside quantize: sites whose TensorRT neighbours cannot run INT8 (e.g., around scaled-dot-product attention) are settled back to FP16, so the exported graph carries Q/DQ only where TensorRT will actually use it.

The quantized graph is re-exported with torch.export and saved to ONNX with external data (the model is >2 GB, above protobuf’s limit).

def export_onnx_external(
    module: nn.Module, args: tuple[torch.Tensor, ...], onnx_path: Path
) -> None:
    """Dynamo-export to ONNX with all tensors in one external .data file.

    :param module:
        The module to export.
    :param args:
        Example positional inputs for the export.
    :param onnx_path:
        Destination path; the tensor data lands next to it in
        ``<name>.data``.
    """
    tmp = onnx_path.with_suffix(".tmp.onnx")
    with torch.no_grad():
        torch.onnx.export(
            module,
            args,
            str(tmp),
            input_names=["image"],
            output_names=["features"],
            do_constant_folding=True,
            dynamo=True,
        )
    m = onnx.load(str(tmp), load_external_data=True)
    onnx.save_model(
        m,
        str(onnx_path),
        save_as_external_data=True,
        all_tensors_to_one_file=True,
        location=onnx_path.name + ".data",
        size_threshold=1024,
    )
    tmp.unlink(missing_ok=True)
    for p in onnx_path.parent.glob(tmp.name + "*"):
        p.unlink(missing_ok=True)


if not ENC_ONNX.exists():
    print("\n[stage 2] Embedl-Deploy transform + INT8 PTQ ...")

    dummy = torch.randn(1, 3, IMG_SIZE, IMG_SIZE, device="cuda")
    graph = torch.export.load(str(FP32_PT2)).module()
    fused = transform(graph, (dummy,), TENSORRT_PATTERNS).model.eval()

    # the only 3-channel conv in the graph is the 16x16 patch embed
    patch_embed = next(
        m
        for m in fused.modules()
        if isinstance(m, nn.Conv2d) and m.in_channels == 3
    )
    quant_cfg = QuantConfig(
        activation=TensorQuantConfig(  # INT8, per-tensor activations
            symmetric=True,
            per_channel=False,
            calibration_method=CalibrationMethod.MINMAX,
        ),
        weight=TensorQuantConfig(  # INT8, per-channel weights
            symmetric=True,
            per_channel=True,
            calibration_method=CalibrationMethod.MINMAX,
        ),
        skip=ModulesToSkip(
            stub={patch_embed},
            weight={patch_embed},
            smooth={nn.LayerNorm},
        ),
    )

    real = torch.from_numpy(enc_input).cuda()
    calib = [real + 0.02 * torch.randn_like(real) for _ in range(N_CALIB)]

    def calibrate(m: nn.Module) -> None:
        """Run the calibration batches through the model."""
        m.eval()
        with torch.no_grad():
            for c in calib:
                m(c)

    print(f"    calibrating on {N_CALIB} jittered copies of the real crop ...")
    qmodel = quantize(
        fused,
        (dummy,),
        config=quant_cfg,
        forward_loop=calibrate,
        freeze_weights=True,
    )

    # fidelity of the QDQ graph itself (before TensorRT enters the picture)
    with torch.no_grad():
        q_feats = qmodel(real).float().cpu().numpy()
    ref = np.load(REF_FEATURES)
    rel = np.abs(q_feats - ref).mean() / (np.abs(ref).mean() + 1e-9)
    STATS["encoder_qdq_rel_diff"] = float(rel)
    print(f"    QDQ-vs-FP32 feature rel diff: {rel * 100:.3f}%")

    print("    exporting INT8 ONNX (external data) ...")
    with torch.no_grad():
        q_exported = torch.export.export(qmodel, (dummy,), strict=False)
    export_onnx_external(q_exported.module(), (dummy,), ENC_ONNX)
    sz = (
        ENC_ONNX.stat().st_size
        + (ART / (ENC_ONNX.name + ".data")).stat().st_size
    )
    print(f"    saved {ENC_ONNX.name} (+ .data, {sz / 1e9:.2f} GB)")
    del graph, fused, qmodel, q_exported
    torch.cuda.empty_cache()
else:
    print("\n[stage 2] reusing cached INT8 encoder ONNX")

Stage 3 — making the decoder + MHR head exportable

The decoder (no-mask prompt add, ray-conditioned camera encoder, 6 promptable decoder layers with per-layer MHR head + camera projection) is wrapped with pure tensor I/O, but cannot be exported as-is:

1. The MHR rig is TorchScript. torch.jit.load-ed modules can’t be traced through. We run the rig once under make_fx(functionalize(...)) to record the dispatched aten ops into a flat FX graph, then clean up the artifacts of that flattening (constant sparse tensors -> dense, in-place copies -> broadcasts, index_copy -> scatter with baked indices).

2. Advanced-index assignment (x[:, idx] = v, boolean-mask writes) exports through a broken index_put lowering. Every such site is replaced by a numerically-identical gather/scatter/where formulation.

3. roma’s rotmat->euler handles gimbal lock with data-dependent nonzero() indexing, which bakes trace-time element counts into the graph. A branchless torch.where version replaces it.

4. Antialiased ray downsampling inside the CameraEncoder — TensorRT’s Resize has no antialias, so rays are downsampled on the host (stage 0) and passed in at feature resolution.

The machinery that fixes all four — FX graph surgery, ONNX-safe method patches, and custom exporter symbolics — lives in the companion module sam_3d_body_export.py. Each patch is numerically identical to the code it replaces (the rig flattening self-checks on five probe input sets before it is accepted), so here we only apply it and export.

class BodyDecoderWrapper(nn.Module):
    """No-mask prompt add + camera encoder + 6-layer promptable decoder (with
    per-layer MHR head, camera projection, and keypoint-token updates) -> final
    pose/mesh outputs. Batch size 1, one dummy keypoint prompt (label -2), no
    mask — exactly the body-only inference configuration.
    """

    def __init__(self, model: nn.Module) -> None:
        """Wrap the full SAM3DBody `model`."""
        super().__init__()
        self.model = model

    def forward(  # noqa: PLR0913 -- one argument per ONNX input
        self,
        features: torch.Tensor,
        rays_ds: torch.Tensor,
        condition_info: torch.Tensor,
        keypoints: torch.Tensor,
        bbox_center: torch.Tensor,
        bbox_scale: torch.Tensor,
        ori_img_size: torch.Tensor,
        cam_int: torch.Tensor,
        affine_trans: torch.Tensor,
        img_size: torch.Tensor,
    ) -> tuple[torch.Tensor, ...]:
        """Run the decoder on one person crop. Input shapes::

        features: (1, 1280, 32, 32) backbone output rays_ds: (1, 2, 32, 32)
        downsampled camera rays condition_info: (1, 3) CLIFF condition
        keypoints: (1, 1, 3) prompt (dummy: [0, 0, -2]) bbox_center: (1, 1, 2)
        bbox_scale: (1, 1, 2) ori_img_size: (1, 1, 2) cam_int: (1, 3, 3)
        affine_trans: (1, 1, 2, 3) img_size: (1, 1, 2)
        """
        m = self.model
        m._batch_size, m._max_num_person = 1, 1  # type: ignore[assignment]  # noqa: SLF001
        m.body_batch_idx, m.hand_batch_idx = [0], []  # type: ignore[assignment]

        # use_mask=False path of _get_mask_prompt reduces to a constant add
        no_mask = m.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1)  # type: ignore[operator, union-attr]
        features = features + no_mask

        batch = {
            # only img.shape[1] (num_person) is read
            "img": features.new_zeros(1, 1, 1),
            "ray_cond": rays_ds,
            "bbox_center": bbox_center,
            "bbox_scale": bbox_scale,
            "ori_img_size": ori_img_size,
            "cam_int": cam_int,
            "affine_trans": affine_trans,
            "img_size": img_size,
        }
        _tokens, pose_outputs = m.forward_decoder(  # type: ignore[operator]
            features,
            init_estimate=None,
            keypoints=keypoints,
            prev_estimate=None,
            condition_info=condition_info,
            batch=batch,
        )
        po = pose_outputs[-1]
        return tuple(po[k] for k in DEC_OUTPUT_NAMES)


def apply_export_patches(model: nn.Module) -> None:
    """Swap every export-hostile site for an ONNX-safe equivalent.

    Every patch is numerically identical to the code it replaces for the
    body-only inference configuration the tutorial exports.

    :param model:
        The ``SAM3DBody`` model, patched in place.
    """
    # rebind in both namespaces: mhr_head imports these by name, so patching
    # the mhr_utils module attribute alone would not reach its call sites
    for mod in (MU, MH):
        mod.compact_cont_to_model_params_body = (
            _safe_compact_cont_to_model_params_body
        )
        mod.compact_cont_to_model_params_hand = (
            _safe_compact_cont_to_model_params_hand
        )
    model.ray_cond_emb.forward = types.MethodType(  # type: ignore[union-attr]
        _camera_encoder_fwd_nodownsample, model.ray_cond_emb
    )
    model.head_pose.replace_hands_in_pose = types.MethodType(  # type: ignore[union-attr]
        _patched_replace_hands_in_pose, model.head_pose
    )
    model.head_pose.forward = types.MethodType(  # type: ignore[union-attr]
        _patched_mhr_head_forward, model.head_pose
    )
    model.head_camera.perspective_projection = types.MethodType(  # type: ignore[union-attr]
        _patched_perspective_projection, model.head_camera
    )
    model.prompt_encoder._embed_keypoints = types.MethodType(  # type: ignore[union-attr]  # noqa: SLF001
        _patched_embed_keypoints, model.prompt_encoder
    )


if not (DEC_ONNX.exists() and FACES_NPY.exists() and REF_MESH.exists()):
    print("\n[stage 3] exporting the decoder + MHR head to ONNX ...")

    # everything on CPU: the flattened rig bakes its trace device into the
    # graph, and the ONNX constant folder needs single-device constants
    body_model, _cfg = load_sam_3d_body(
        CKPT_PATH, device="cpu", mhr_path=MHR_PATH
    )
    np.save(FACES_NPY, body_model.head_pose.faces.cpu().numpy())
    apply_export_patches(body_model)
    wrapper = BodyDecoderWrapper(body_model).eval()

    ref_features = np.load(REF_FEATURES)
    dec_args = tuple(
        torch.from_numpy({"features": ref_features, **dec_feeds}[k])
        for k in DEC_INPUT_NAMES
    )

    print("    flattening the TorchScript MHR rig to FX ...")
    flatten_rig_for_export(wrapper, body_model.head_pose, MHR_PATH, dec_args)

    # FP32 eager reference mesh for the real sample -- the ground truth that
    # the TensorRT pipeline is compared against in stage 5
    with torch.no_grad():
        ref_out = dict(zip(DEC_OUTPUT_NAMES, wrapper(*dec_args), strict=True))
    np.savez(
        REF_MESH,
        vertices=ref_out["pred_vertices"][0].numpy(),
        cam_t=ref_out["pred_cam_t"][0].numpy(),
    )
    print(
        f"    FP32 reference mesh: {ref_out['pred_vertices'].shape[1]} "
        f"vertices"
    )

    register_ts_exporter_symbolics()
    print("    exporting ONNX (TorchScript tracer, CPU) ...")
    with torch.no_grad():
        torch.onnx.export(
            wrapper,
            dec_args,
            str(DEC_ONNX),
            input_names=DEC_INPUT_NAMES,
            output_names=DEC_OUTPUT_NAMES,
            opset_version=17,
            dynamo=False,
            do_constant_folding=True,
        )
    print(
        f"    saved {DEC_ONNX.name} ({DEC_ONNX.stat().st_size / 1e6:.0f} MB)"
    )
    del body_model, wrapper
else:
    print("\n[stage 3] reusing cached decoder ONNX + reference mesh")

Stage 4 — TensorRT engines

TensorRT 11 builds strongly-typed networks: every layer keeps the precision declared in the ONNX file, and the old FP16/INT8 builder flags are gone. So we set precisions in the ONNX itself:

• FP32 tensors, Cast targets, and inline constants are rewritten to FP16 (clipped to the FP16 range). The encoder’s Q/DQ INT8 structure is untouched — those layers run INT8, the rest FP16.

• The decoder’s float64 mesh-rig section (FK chains accumulate error) is demoted to FP32, not FP16, to keep the skinning math accurate.

The cast and the engine build each happen once and are cached on disk.

TRT_LOG = trt.Logger(trt.Logger.WARNING)
ENC_ENGINE = ART / "embedl_sam3dbody_int8.engine"
DEC_ENGINE = ART / "sam3dbody_decoder.engine"

ELEM_REMAP = {
    TensorProto.FLOAT: TensorProto.FLOAT16,
    TensorProto.DOUBLE: TensorProto.FLOAT,
}


def _retype_tensor(t: onnx.TensorProto) -> None:
    """Rewrite a tensor fp32 -> clipped fp16 or fp64 -> fp32, in place."""
    if t.data_type == TensorProto.FLOAT:
        arr = np.clip(numpy_helper.to_array(t), -65504, 65504).astype(
            np.float16
        )
        t.CopyFrom(numpy_helper.from_array(arr, t.name))
    elif t.data_type == TensorProto.DOUBLE:
        arr = numpy_helper.to_array(t).astype(np.float32)
        t.CopyFrom(numpy_helper.from_array(arr, t.name))


def cast_onnx_to_fp16(onnx_path: Path) -> Path:
    """Cast an fp32 ONNX file to fp16 for a strongly-typed TensorRT build.

    fp64 becomes fp32 and fp32 becomes fp16 across initializers, value infos,
    Cast targets, and inline tensor attributes; the Q/DQ INT8 structure is
    untouched.

    :param onnx_path:
        The fp32 ONNX file to cast.
    :returns:
        The path of the cast copy (cached across runs).
    """
    fp16_path = onnx_path.with_name(onnx_path.stem + "_fp16.onnx")
    if fp16_path.exists():
        return fp16_path
    print(f"    casting {onnx_path.name} to fp16 (one-time) ...")
    model = onnx.load(str(onnx_path))
    g = model.graph
    for init in g.initializer:
        _retype_tensor(init)
    for vi in list(g.input) + list(g.output) + list(g.value_info):
        t = vi.type.tensor_type
        t.elem_type = ELEM_REMAP.get(t.elem_type, t.elem_type)
    for node in g.node:
        for a in node.attribute:
            if node.op_type == "Cast" and a.name == "to":
                a.i = ELEM_REMAP.get(a.i, a.i)
            elif a.type == onnx.AttributeProto.TENSOR:
                _retype_tensor(a.t)
    onnx.save(
        model,
        str(fp16_path),
        save_as_external_data=True,
        location=fp16_path.name + ".data",
        all_tensors_to_one_file=True,
    )
    return fp16_path


def build_engine(onnx_path: Path, engine_path: Path) -> trt.ICudaEngine:
    """Build (or load cached) a strongly-typed fp16 engine from fp32 ONNX.

    :param onnx_path:
        The fp32 ONNX file (cast to fp16 on first use).
    :param engine_path:
        Where the serialized engine is cached.
    :returns:
        The deserialized TensorRT engine.
    """
    trt.init_libnvinfer_plugins(TRT_LOG, "")
    if engine_path.exists():
        return trt.Runtime(TRT_LOG).deserialize_cuda_engine(
            engine_path.read_bytes()
        )
    onnx_path = cast_onnx_to_fp16(onnx_path)
    print(f"    building {engine_path.name} from {onnx_path.name} ...")
    builder = trt.Builder(TRT_LOG)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
    )
    parser = trt.OnnxParser(network, TRT_LOG)
    if not parser.parse_from_file(str(onnx_path)):
        raise RuntimeError(
            "\n".join(
                str(parser.get_error(i)) for i in range(parser.num_errors)
            )
        )
    cfg = builder.create_builder_config()
    cfg.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 * 1024**3)
    t0 = time.perf_counter()
    plan = builder.build_serialized_network(network, cfg)
    if plan is None:
        raise RuntimeError("engine build failed")
    engine_path.write_bytes(bytes(plan))
    print(
        f"    {engine_path.name}: built in {time.perf_counter() - t0:.0f}s, "
        f"{engine_path.stat().st_size / 1e6:.0f} MB"
    )
    return trt.Runtime(TRT_LOG).deserialize_cuda_engine(bytes(plan))

Stage 5 — the TensorRT demo: real image in, 3D mesh out

A small pycuda runner drives both engines: pinned host buffers, one device buffer per I/O tensor, execute_async_v3 on a private stream. Latency is measured on engine execution only (inputs already on device), which is the number that matters when the pipeline runs inside a larger system.

class Engine:
    """Single-context TensorRT runner with engine-only latency measurement."""

    def __init__(self, onnx_path: Path, engine_path: Path) -> None:
        """Build or load the engine and allocate its I/O buffers."""
        self.engine = build_engine(onnx_path, engine_path)
        self.ctx = self.engine.create_execution_context()
        self.stream = cuda.Stream()
        self.host: dict[str, npt.NDArray[np.generic]] = {}
        self.dev: dict[str, cuda.DeviceAllocation] = {}
        self.shapes: dict[str, tuple[int, ...]] = {}
        self.inputs: list[str] = []
        self.outputs: list[str] = []
        for i in range(self.engine.num_io_tensors):
            nm = self.engine.get_tensor_name(i)
            shape = tuple(self.engine.get_tensor_shape(nm))
            dt = trt.nptype(self.engine.get_tensor_dtype(nm))
            self.host[nm] = cuda.pagelocked_empty(int(np.prod(shape)) or 1, dt)
            self.dev[nm] = cuda.mem_alloc(max(self.host[nm].nbytes, 8))
            self.shapes[nm] = shape
            self.ctx.set_tensor_address(nm, int(self.dev[nm]))
            is_in = self.engine.get_tensor_mode(nm) == trt.TensorIOMode.INPUT
            (self.inputs if is_in else self.outputs).append(nm)

    def __call__(
        self, feeds: dict[str, npt.NDArray[np.float32]]
    ) -> dict[str, npt.NDArray[np.float32]]:
        """Run one inference on `feeds` and return the fp32 outputs."""
        for nm in self.inputs:
            np.copyto(
                self.host[nm], feeds[nm].ravel().astype(self.host[nm].dtype)
            )
            cuda.memcpy_htod_async(self.dev[nm], self.host[nm], self.stream)
        self.ctx.execute_async_v3(self.stream.handle)
        for nm in self.outputs:
            cuda.memcpy_dtoh_async(self.host[nm], self.dev[nm], self.stream)
        self.stream.synchronize()
        return {
            nm: self.host[nm]
            .reshape(self.shapes[nm])
            .astype(np.float32)
            .copy()
            for nm in self.outputs
        }

    def benchmark(self) -> tuple[float, float]:
        """Mean and p95 engine-only latency in milliseconds."""

        def _run() -> None:
            self.ctx.execute_async_v3(self.stream.handle)
            self.stream.synchronize()

        return benchmark_ms(_run)


print("\n[stage 4] TensorRT engines ...")
encoder = Engine(ENC_ONNX, ENC_ENGINE)
decoder = Engine(DEC_ONNX, DEC_ENGINE)

print("\n[stage 5] running the sample end to end on TensorRT ...")
features = encoder({"image": enc_input})["features"]
out = decoder({**dec_feeds, "features": features})
V = out["pred_vertices"][0]
cam_t = out["pred_cam_t"][0]
print(f"    recovered mesh: {V.shape[0]} vertices")

# fidelity vs the eager FP32 reference
ref_feats = np.load(REF_FEATURES)
feat_rel = np.abs(features - ref_feats).mean() / (
    np.abs(ref_feats).mean() + 1e-9
)
ref_mesh = np.load(REF_MESH)
mesh_d = np.abs(V - ref_mesh["vertices"])
STATS.update(
    {
        "encoder_trt_rel_diff": float(feat_rel),
        "mesh_max_mm": float(mesh_d.max() * 1000),
        "mesh_mean_mm": float(mesh_d.mean() * 1000),
    }
)
print(
    f"    encoder features INT8-TRT vs FP32 eager: "
    f"rel diff {feat_rel * 100:.3f}%"
)
print(
    f"    full-pipeline mesh vs FP32 eager: "
    f"max {mesh_d.max() * 1000:.2f} mm  mean {mesh_d.mean() * 1000:.3f} mm"
)

enc_ms, enc_p95 = encoder.benchmark()
dec_ms, dec_p95 = decoder.benchmark()
STATS.update(
    {
        "encoder_trt_ms": enc_ms,
        "encoder_trt_p95_ms": enc_p95,
        "decoder_trt_ms": dec_ms,
        "decoder_trt_p95_ms": dec_p95,
        "tensorrt": trt.__version__,
        "torch": torch.__version__,
        "gpu": torch.cuda.get_device_name(0),
        "input_size": IMG_SIZE,
    }
)

Rendering — mesh overlay + turntable

The mesh is projected back onto the crop with the same intrinsics + affine used in preprocessing, flat-shaded, and painter-sorted (no GL dependency).

LIGHT = np.array([0.3, 0.5, 1.0])
LIGHT /= np.linalg.norm(LIGHT)
SKIN = np.array([0.80, 0.78, 0.72])


def _shade(
    v: npt.NDArray[np.float32], f: npt.NDArray[np.int64]
) -> npt.NDArray[np.float64]:
    """Per-face flat-shading colors from a fixed directional light."""
    n = np.cross(v[f][:, 1] - v[f][:, 0], v[f][:, 2] - v[f][:, 0])
    n /= np.linalg.norm(n, axis=1, keepdims=True) + 1e-9
    lam = np.clip(np.abs(n @ LIGHT), 0, 1)[:, None]
    return np.clip(0.25 + 0.75 * lam * SKIN, 0, 1)


def _view(
    ax: Axes,
    verts: npt.NDArray[np.float32],
    faces: npt.NDArray[np.int64],
    deg: float,
    title: str,
) -> None:
    """Draw the mesh into `ax`, rotated `deg` degrees about the y axis."""
    Vc = verts - verts.mean(0)
    th = np.radians(deg)
    R = np.array(
        [[np.cos(th), 0, np.sin(th)], [0, 1, 0], [-np.sin(th), 0, np.cos(th)]]
    )
    Vr = Vc @ R.T
    p = Vr[:, :2] * [1, -1]
    o = np.argsort(Vr[faces].mean(1)[:, 2])
    ax.add_collection(
        PolyCollection(
            p[faces][o], facecolors=_shade(Vr, faces)[o], edgecolors="none"
        )
    )
    ax.set_xlim(p[:, 0].min(), p[:, 0].max())
    ax.set_ylim(p[:, 1].min(), p[:, 1].max())
    ax.set_aspect("equal")
    ax.axis("off")
    ax.set_title(title, fontsize=11)


def render_overlay(
    image_path: Path,
    verts: npt.NDArray[np.float32],
    faces: npt.NDArray[np.int64],
    cam_t: npt.NDArray[np.float32],
    crop: npt.NDArray[np.uint8],
    cam_int: npt.NDArray[np.float32],
    affine: npt.NDArray[np.float32],
    out: Path,
) -> None:
    """Save a 4-panel figure: input, crop overlay, and two mesh views.

    :param image_path:
        Path of the original input image.
    :param verts:
        Mesh vertices, shape ``(N, 3)``.
    :param faces:
        Mesh faces, shape ``(M, 3)``.
    :param cam_t:
        Predicted camera translation, shape ``(3,)``.
    :param crop:
        The 512x512 RGB model crop from preprocessing.
    :param cam_int:
        Camera intrinsics used in preprocessing, shape ``(1, 3, 3)``.
    :param affine:
        The full-image-to-crop affine from preprocessing.
    :param out:
        Destination PNG path.
    """
    img = read_rgb(image_path)
    fig, ax = plt.subplots(1, 4, figsize=(15, 6))
    fig.patch.set_facecolor("white")
    ax[0].imshow(img)
    ax[0].axis("off")
    ax[0].set_title("Input", fontsize=11)
    # project with the camera intrinsics, then map full-image pixels into the
    # model crop with the same affine used for preprocessing
    Vc = verts + cam_t
    z = np.clip(Vc[:, 2], 1e-3, None)
    f, cx, cy = cam_int[0, 0, 0], cam_int[0, 0, 2], cam_int[0, 1, 2]
    pf = np.stack(
        [f * Vc[:, 0] / z + cx, f * Vc[:, 1] / z + cy, np.ones_like(z)], 1
    )
    A = affine.reshape(2, 3).astype(np.float64)
    p = pf @ A.T
    o = np.argsort(-Vc[faces].mean(1)[:, 2])
    ax[1].imshow(crop)
    ax[1].add_collection(
        PolyCollection(
            p[faces][o],
            facecolors=_shade(Vc, faces)[o],
            edgecolors="none",
            alpha=0.8,
        )
    )
    ax[1].set_xlim(0, IMG_SIZE)
    ax[1].set_ylim(IMG_SIZE, 0)
    ax[1].axis("off")
    ax[1].set_title("Mesh overlay", fontsize=11)
    _view(ax[2], verts, faces, 20, "¾ view")
    _view(ax[3], verts, faces, 90, "side view")
    plt.tight_layout()
    plt.savefig(out, dpi=160, bbox_inches="tight")
    plt.close()
    print(f"    wrote {out}")


def render_turntable(
    verts: npt.NDArray[np.float32],
    faces: npt.NDArray[np.int64],
    out: Path,
) -> None:
    """Save a rotating-mesh GIF (one frame every 15 degrees).

    :param verts:
        Mesh vertices, shape ``(N, 3)``.
    :param faces:
        Mesh faces, shape ``(M, 3)``.
    :param out:
        Destination GIF path.
    """
    frames = []
    for a in range(0, 360, 15):
        fig, ax = plt.subplots(figsize=(4, 6))
        fig.patch.set_facecolor("white")
        _view(ax, verts, faces, a, "")
        fig.canvas.draw()
        rgba = fig.canvas.buffer_rgba()
        frames.append(np.asarray(rgba)[..., :3].copy())
        plt.close()
    imageio.mimsave(out, frames, duration=0.1, loop=0)
    print(f"    wrote {out}")


faces = np.load(FACES_NPY)
render_overlay(
    SAMPLE_IMG,
    V,
    faces,
    cam_t,
    crop_rgb,
    dec_feeds["cam_int"],
    dec_feeds["affine_trans"],
    ART / "mesh_demo_trt.png",
)
render_turntable(V, faces, ART / "mesh_demo_trt_spin.gif")

Summary

STATS_PATH.write_text(json.dumps(STATS, indent=2))
eager = STATS.get("encoder_fp32_eager_ms")
eager_ms = eager if isinstance(eager, float) else None
print("\n" + "=" * 72)
print(
    f"SAM 3D Body on TensorRT {trt.__version__} -- "
    f"{torch.cuda.get_device_name(0)}, {IMG_SIZE}x{IMG_SIZE}"
)
print("=" * 72)
print(f"{'Stage':<28}{'Mean':>10}{'p95':>10}{'Engine':>12}")
print("-" * 72)
if eager_ms:
    print(
        f"{'encoder FP32 (eager torch)':<28}{eager_ms:>7.2f} ms"
        f"{'--':>10}{'--':>12}"
    )
for name, ms, p95, engine_path in (
    ("encoder INT8+FP16 (TRT)", enc_ms, enc_p95, ENC_ENGINE),
    ("decoder FP16 (TRT)", dec_ms, dec_p95, DEC_ENGINE),
):
    print(
        f"{name:<28}{ms:>7.2f} ms{p95:>7.2f} ms"
        f"{engine_path.stat().st_size / 2**20:>8.0f} MiB"
    )
print(f"{'total (TRT)':<28}{enc_ms + dec_ms:>7.2f} ms")
if eager_ms:
    print(f"\nencoder speedup vs eager FP32: {eager_ms / enc_ms:.2f}x")
print(
    f"encoder features rel diff (INT8 engine vs FP32 eager): "
    f"{STATS['encoder_trt_rel_diff'] * 100:.3f}%"
)
print(
    f"mesh error vs FP32 eager pipeline: mean "
    f"{STATS['mesh_mean_mm']:.3f} mm, max {STATS['mesh_max_mm']:.2f} mm"
)
print("=" * 72)
print(f"\nresults written to {STATS_PATH}")