AXDocLayout-M v1.2 — ONNX CUDA + TF32
CUDA-powered document layout detection. AX DocLayout locates layout regions
in a decoded RGB page image and returns class labels, confidence scores and
bounding boxes in the original raster's pixel coordinates.
What changed in v1.2
This is a CUDA runtime release using the same FP32 weights and ONNX graph as
v1.1. The ONNX backend now explicitly enables ONNX Runtime's CUDA execution
provider and TF32. Postprocessing is vectorized with equivalent raster
geometry. The source package verifies the CUDA provider, effective TF32 policy
and placement of major compute operations. CUDA setup failures raise an error;
the loader does not retry with a CPU-only session.
The architecture remains RT-DETRv2 rtdetrv2_r50vd_m, with 11 layout classes.
There is no integer quantization or model retraining in this release. Both
model.onnx (standard operators, opset 16) and FP32 model.safetensors are
provided. Complete inference source is in inference/; attribution is in NOTICE.
CUDA requirements
Use Linux x86-64, Python 3.12 or 3.13, a supported NVIDIA GPU with TF32 capability
(Ampere or newer), and a driver compatible with CUDA 12.8. The ONNX extra pins
ONNX Runtime GPU 1.23.2 and includes CUDA/cuDNN Python dependencies. The tested
configuration is RTX 5090, CUDA 12.8 and cuDNN 9.10.2.21. The ONNX path requires
CUDA. backend="pytorch" remains an explicit CPU FP32 comparison path and does
not provide the CUDA performance below.
Floating-point CUDA execution can produce small output differences between
invocations or sessions. Outputs are not promised to be bitwise reproducible.
TF32 is an execution-provider setting. ONNX inputs and stored model weights
remain FP32; downloading the ONNX file alone does not enable TF32.
Reference development results
The following aggregate results are from internal development validation,
not an independent held-out acceptance test. These are measured results for
one hardware/runtime configuration, not a latency guarantee.
| Runtime configuration |
p50 ms/page |
p95 ms/page |
Serial pages/s |
AP |
AP50 |
AP70 |
AP90 |
| Prior eager PyTorch FP32 reference |
46.242 |
62.540 |
19.48 |
76.2382 |
90.9897 |
86.8844 |
55.8986 |
| Optimized ONNX CUDA FP32 |
34.423 |
40.770 |
27.76 |
76.2384 |
90.9896 |
86.8844 |
55.8988 |
| Optimized ONNX CUDA + TF32 |
25.792 |
32.395 |
35.91 |
76.2414 |
90.9891 |
86.8792 |
55.9025 |
AP uses all 11 classes on 6,489 development pages, with a 0.001 score floor
for evaluation. The public prediction API uses a 0.30 score floor.
Latency uses batch one,
three counterbalanced passes of 100 pages after untimed startup samples.
The timed region is decoded RGB to CPU box rows, including letterboxing,
normalization, transfers, inference, validation and postprocessing; it excludes
image-file decoding, IPC, initialization and evaluation. API serialization
costs and host-load variation can change application latency. The measured p50
reduction is 44.2% versus the eager reference. Small AP differences do not
establish an accuracy improvement.
On two additional financial-document development cohorts (256 and 100 pages),
table AP changed by -0.0639 and +0.0730 points. The largest class AP decrease
on the general development cohort was 0.0328 points. Corpus and per-page data
are not included.
The tested graph placed convolutions, matrix multiplications, GridSample and
Softmax on CUDA. Some shape/control and arithmetic nodes execute on CPU;
this is not an all-GPU graph claim.
Release installation validation
A separate clean installation of the CUDA-only inference package, with no
PyTorch installed, was evaluated on all 6,489 general development pages:
| AP |
AP50 |
AP70 |
AP90 |
| 76.2396 |
90.9855 |
86.8767 |
55.9099 |
The same installation passed the predefined development quality margins on
both financial-document cohorts. An additional 128-page API check verified
preprocessing and exact postprocessing of the captured inference outputs.
Independent CUDA invocations are not bitwise-equivalent; these checks do not
claim cross-session raw-detection equivalence or independent acceptance.
Input and output
Supply one HWC uint8 RGB page raster. Preprocessing preserves aspect ratio,
resizes the long edge to 1024, pads to 1024 x 1024 with white and normalizes to
NCHW float32 in [0, 1], under rgb-letterbox-1024/1. The Python API applies
score floor 0.30, removes reserved/invalid labels, reverses letterboxing and
clips boxes to the original raster. Index 0 is reserved.
The JSON schema is ax-doclayout.detection/1, taxonomy ax-doclayout-11/1.
Boxes use half-open pixel edges with top-left origin, x rightward and y
downward. Raw ONNX inputs are images: float32 [N, 3, 1024, 1024] and
orig_target_sizes: int64 [N, 2], [1024, 1024] per image. Raw outputs are
labels, boxes (letterbox-space xyxy) and scores.
| ID |
Label |
| 1 |
caption |
| 2 |
footnote |
| 3 |
formula |
| 4 |
list_item |
| 5 |
page_footer |
| 6 |
page_header |
| 7 |
picture |
| 8 |
section_header |
| 9 |
table |
| 10 |
text |
| 11 |
title |
Use and limitations
Validate the pinned release on representative pages before deployment.
Boundary precision, very small regions and unusual document designs can need
review. Confidence scores do not guarantee correctness. The model accepts
page rasters; it does not open PDFs, perform OCR, extract text, predict reading
order/hierarchy, or reconstruct table cells. Larger-batch throughput and other
GPU/software configurations were not qualified by these results.
License
The weights and modified inference source are AGPL-3.0-only. Commercial use
is permitted subject to the license terms. See LICENSE and NOTICE for terms,
source-availability requirements and preserved Apache-2.0 component notices.
This does not relicense previous releases. Complete inference source is supplied
in inference/; datasets and the training pipeline are not distributed.
Installation and inference
Use the immutable source and weight revisions in the root README quick-start.
Install the ONNX GPU extra and explicitly select backend="onnx" for CUDA TF32.
An explicit CPU comparison is available through inference[torch] and
backend="pytorch". The package verifies payload digests before inference.
Runtime reference: https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html
Quick start — CUDA + TF32
Requires supported Linux x86-64/NVIDIA CUDA setup and Python 3.12 or 3.13.
Install the download helper, then the immutable inference source:
python -m pip install "huggingface-hub>=0.36,<2"
import subprocess
import sys
from huggingface_hub import snapshot_download
source = snapshot_download(
"AutomatosX/AXDocLayout-M-v1.2",
revision="1b4ed073124fddb0c8dce4f915b7b22aa9fab8ee",
allow_patterns=["inference/**"],
token=False,
)
subprocess.check_call([
sys.executable, "-m", "pip", "install", source + "/inference[onnx]"
])
Run CUDA-powered layout detection:
import numpy as np
from PIL import Image
from ax_doclayout import AXDocLayout, PageImage
model = AXDocLayout.from_pretrained(
"AutomatosX/AXDocLayout-M-v1.2",
revision="7efd000fc7a1ef50ed7913835f2d8f10f30dc8e3",
backend="onnx", # v1.2: CUDAExecutionProvider with TF32 enabled
)
page = PageImage(
rgb=np.asarray(Image.open("page.png").convert("RGB"), dtype=np.uint8),
page_index=0,
)
result = model.predict(page)
print(result.to_json())
For the separate CPU PyTorch comparison, install inference[torch] and
explicitly select backend="pytorch". The CUDA ONNX loader does not switch
to this backend automatically.
