# Copyright (C) 2026 Embedl AB

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

`SAM 3D Body <https://github.com/facebookresearch/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 :download:`sam_3d_body_export.py
</tutorials/sam_3d_body/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}")