export_pp_doclayout_v3.py

"""Export PaddlePaddle/PP-DocLayoutV3_safetensors to ONNX.

The HF module returns a large structured output (intermediate decoder states,
per-layer masks, reference points...). Only four tensors are needed at serving
time, so we wrap the model and return exactly those:

    logits        (B, 300, num_classes)
    pred_boxes    (B, 300, 4)          cxcywh, normalised
    order_logits  (B, 300, 300)        reading-order pointer matrix
    out_masks     (B, 300, 200, 200)   mask logits @ stride 4  (optional)

`out_masks` is by far the biggest tensor (300*200*200*4B = 48 MB per image in
fp32). Export with --no-masks if you only need boxes + reading order; polygons
then degrade to axis-aligned rectangles.

Usage:
    python export_pp_doclayout_v3.py --output pp_doclayoutv3.onnx
    python export_pp_doclayout_v3.py --output pp_doclayoutv3_nomask.onnx --no-masks
    python export_pp_doclayout_v3.py --output m.onnx --fp16   # extra fp16 copy
"""

from __future__ import annotations

import argparse
from contextlib import contextmanager
from pathlib import Path

import numpy as np
import torch
from torch import nn
from transformers import AutoModelForObjectDetection

DEFAULT_MODEL = "PaddlePaddle/PP-DocLayoutV3_safetensors"
INPUT_SIZE = 800  # preprocessor_config.json: {"height": 800, "width": 800}


@contextmanager
def fp32_position_embedding():
    """Force the 2D sin/cos position embedding to float32 during tracing.

    Upstream computes the frequency grid in float64, which bakes Sin/Cos(double)
    nodes into the graph. ONNX Runtime's CPU EP has no double kernel for those, so
    the session fails to load with:
        NOT_IMPLEMENTED : Could not find an implementation for Cos(7)
    The embedding is cast to float32 anyway; the difference is ~1e-6.
    """
    from transformers.models.pp_doclayout_v3 import modeling_pp_doclayout_v3 as M

    original = M.build_2d_sinusoidal_position_embedding

    def patched(height, width, embed_dim=256, temperature=10000.0,
                cls_token=False, device=None, dtype=torch.float32):
        pos_dim = embed_dim // 4
        omega = torch.arange(pos_dim, dtype=torch.float32, device=device) / pos_dim
        omega = 1.0 / temperature**omega
        grid_h = torch.arange(height, dtype=torch.float32, device=device)
        grid_w = torch.arange(width, dtype=torch.float32, device=device)
        grid_h, grid_w = torch.meshgrid(grid_h, grid_w, indexing="ij")
        emb_h = grid_h.flatten().outer(omega)
        emb_w = grid_w.flatten().outer(omega)
        pos_embed = torch.cat([emb_h.sin(), emb_h.cos(), emb_w.sin(), emb_w.cos()], dim=1)
        if cls_token:
            zeros = torch.zeros(1, embed_dim, dtype=torch.float32, device=device)
            pos_embed = torch.cat([zeros, pos_embed], dim=0)
        return pos_embed.to(dtype)

    M.build_2d_sinusoidal_position_embedding = patched
    try:
        yield
    finally:
        M.build_2d_sinusoidal_position_embedding = original


class PPDocLayoutV3ExportWrapper(nn.Module):
    """Flattens the HF output struct down to the tensors post-processing needs."""

    def __init__(self, model: nn.Module, with_masks: bool = True) -> None:
        super().__init__()
        self.model = model
        self.with_masks = with_masks

    def forward(self, pixel_values: torch.Tensor):
        out = self.model(pixel_values=pixel_values)

        order_logits = out.order_logits
        # Doc types it as a tuple; take the final decoder layer if so.
        if isinstance(order_logits, (tuple, list)):
            order_logits = order_logits[-1]
        if order_logits.dim() == 4:  # (B, layers, Q, Q)
            order_logits = order_logits[:, -1]

        if not self.with_masks:
            return out.logits, out.pred_boxes, order_logits

        masks = out.out_masks
        if masks.dim() == 5:  # (B, layers, Q, H, W) -> last layer only
            masks = masks[:, -1]
        return out.logits, out.pred_boxes, order_logits, masks


def export(
    model_path: str,
    output: Path,
    with_masks: bool = True,
    opset: int = 17,
    dynamo: bool = False,
    dynamic_batch: bool = True,
) -> tuple[nn.Module, list[str]]:
    model = AutoModelForObjectDetection.from_pretrained(model_path, dtype=torch.float32).eval()
    # Pure-PyTorch deformable attention (grid_sample) instead of the custom CUDA
    # kernel — the custom op has no ONNX symbolic. The checkpoint config already
    # sets this, but be explicit in case someone overrides it.
    model.config.disable_custom_kernels = True

    wrapper = PPDocLayoutV3ExportWrapper(model, with_masks=with_masks).eval()

    dummy = torch.randn(1, 3, INPUT_SIZE, INPUT_SIZE)
    output_names = ["logits", "pred_boxes", "order_logits"]
    if with_masks:
        output_names.append("out_masks")

    # Height/width stay static: the image processor always resizes to 800x800,
    # and static spatial dims let ORT constant-fold the anchor generation.
    dynamic_axes = None
    if dynamic_batch:
        dynamic_axes = {name: {0: "batch"} for name in ["pixel_values"] + output_names}

    output.parent.mkdir(parents=True, exist_ok=True)
    with torch.inference_mode(), fp32_position_embedding():
        torch.onnx.export(
            wrapper,
            (dummy,),
            str(output),
            input_names=["pixel_values"],
            output_names=output_names,
            dynamic_axes=dynamic_axes,
            opset_version=opset,
            do_constant_folding=True,
            dynamo=dynamo,
        )
    print(f"Exported -> {output}  ({output.stat().st_size / 1e6:.1f} MB)")
    return wrapper, output_names


def check_parity(
    wrapper: nn.Module, onnx_path: Path, output_names: list[str], batch: int = 2, atol: float = 1e-3
) -> None:
    """Compare ONNX Runtime against PyTorch on random input.

    Note: this is only meaningful with the real pretrained weights. With randomly
    initialised weights the encoder emits thousands of identical proposal scores,
    so TopK query selection is arbitrary and torch/ORT legitimately disagree.
    """
    import onnxruntime as ort

    x = torch.randn(batch, 3, INPUT_SIZE, INPUT_SIZE)
    with torch.inference_mode():
        torch_out = [t.numpy() for t in wrapper(x)]

    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
    onnx_out = sess.run(None, {"pixel_values": x.numpy()})

    ok = True
    for name, t, o in zip(output_names, torch_out, onnx_out):
        diff = np.abs(t - o).max()
        status = "OK " if diff < atol else "FAIL"
        ok &= diff < atol
        print(f"  {status} {name:14s} shape={tuple(o.shape)} max|diff|={diff:.3e}")
    print("Parity check passed\n" if ok else "Parity check FAILED — do not ship this graph\n")


def to_fp16(onnx_path: Path) -> Path:
    """Half-precision copy. Keep normalisation/mask ops in fp32 to avoid overflow."""
    import onnx
    from onnxconverter_common import float16

    model = onnx.load(str(onnx_path))
    fp16_model = float16.convert_float_to_float16(
        model, keep_io_types=True, op_block_list=["GridSample", "ReduceMean", "Pow", "Sqrt", "Div"]
    )
    out = onnx_path.with_name(onnx_path.stem + "_fp16.onnx")
    onnx.save(fp16_model, str(out))
    print(f"fp16 -> {out}  ({out.stat().st_size / 1e6:.1f} MB)")
    return out


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--model", default=DEFAULT_MODEL)
    p.add_argument("--output", type=Path, default=Path("pp_doclayoutv3.onnx"))
    p.add_argument("--opset", type=int, default=17, help=">=16 required for GridSample")
    p.add_argument("--no-masks", dest="with_masks", action="store_false")
    p.add_argument("--static-batch", dest="dynamic_batch", action="store_false")
    p.add_argument("--dynamo", action="store_true", help="use the TorchDynamo exporter")
    p.add_argument("--fp16", action="store_true", help="also emit an fp16 copy")
    p.add_argument("--skip-check", action="store_true")
    args = p.parse_args()

    wrapper, names = export(
        args.model,
        args.output,
        with_masks=args.with_masks,
        opset=args.opset,
        dynamo=args.dynamo,
        dynamic_batch=args.dynamic_batch,
    )
    if not args.skip_check:
        check_parity(wrapper, args.output, names, batch=2 if args.dynamic_batch else 1)
    if args.fp16:
        to_fp16(args.output)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())