AutomatosX/AXDocLayout-M-v1.2-CUDA-TF32

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.

On the measured RTX 5090, one-page CUDA inference reserved 1048 MiB of GPU memory for this process and peaked at 2471 MiB of resident host RAM. The CUDA graph is 127.1 MiB. The Linux ONNX extra installs 2.94 GiB of runtime libraries. Measurement conditions and larger batches are in Memory.

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.

Memory

Observed on one RTX 5090 (32607 MiB, driver 615.71.09) with Python 3.12.13, ONNX Runtime GPU 1.23.2, CUDA runtime 12.8.90 and cuDNN 9.10.2.21. Stored weights remain FP32. The 36.7M parameters occupy about 140 MiB; the GPU figures below are larger because they include activations and the CUDA workspace. Another GPU, driver or concurrent process can reserve a different amount. Smaller GPUs were not tested. These rows are not a latency or accuracy result.

One-page API

backend="onnx" keeps one CUDA session after a startup probe. The input was a synthetic blank page. Every page is letterboxed to 1024 x 1024, so this GPU reservation does not grow with the source raster's pixel count. nvidia-smi reported 10 MiB in use before the process started. A second process reported the same 784 MiB and 1048 MiB process reservations.

Stage	Process GPU memory	nvidia-smi memory.used
Session loaded	784 MiB	802 MiB
After public predict calls	1048 MiB	1066 MiB

That process peaked at 2471 MiB resident host RAM (2,530,632 KiB) and was still at 2245 MiB after inference. CUDA maps a much larger virtual address range. Virtual size is not RAM that has to be free.

Raw ONNX batch

The graph input images has a dynamic batch N. The public Python API still submits one page. Two fresh sessions, with the published CUDA TF32 options and synthetic 1024 inputs, reserved:

Batch	Higher process GPU memory	Higher nvidia-smi memory.used
1	1054 MiB	1072 MiB
2	2582 MiB	2600 MiB
4	2590 MiB	2608 MiB
8	3614 MiB	3632 MiB

Batch 2 reserved 1566 MiB in one session and 2582 MiB in the other. The CUDA workspace choice is not stable across sessions, so plan for the higher figure. Batch 4 stayed within 8 MiB of that higher batch-2 reservation. These samples do not qualify accuracy or throughput above batch 1.

Disk

Item	Size
model.onnx, the CUDA graph	127.1 MiB (133,276,847 bytes)
model.safetensors, CPU PyTorch comparison only	140.0 MiB (146,766,504 bytes)
Whole repository	267.1 MiB (280,043,351 bytes)
Installed Linux x86-64 inference[onnx] runtime libraries	2.94 GiB (3,154,851,060 bytes)

The CUDA path does not read model.safetensors. The installed library figure is ONNX Runtime GPU 1.23.2 plus the pinned CUDA 12.8 and cuDNN 9.10 packages on this machine. Pip printed about 1.99 GB for those wheel files before they were unpacked.

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 averages IoU thresholds 0.50:0.95; AP50, AP70 and AP90 use their respective IoU thresholds. Scores are percentages. Evaluation definitions are private, so these numbers are not independently reproducible benchmarks. The public prediction API uses a 0.30 score floor.

Latency measures batch-one decoded RGB to CPU box rows, including preprocessing, transfers, inference, validation and postprocessing. It excludes image decoding, IPC, initialization, evaluation and API serialization. Application latency varies with host load. The measured median reduction is 44.2% versus the eager reference. Small AP differences do not establish an accuracy improvement.

Two additional internal evaluations showed AP changes of -0.0639 and +0.0730 points. The largest class AP decrease in the primary evaluation was 0.0328 points. Validate on representative inputs before deployment.

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.

Published inference runtime

The CUDA-only inference package does not require PyTorch. Its internal development accuracy measurements are:

AP	AP50	AP70	AP90
76.2396	90.9855	86.8767	55.9099

These results are not independent acceptance. Independent CUDA invocations are not bitwise equivalent.

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. Observed memory for the one-page API and for raw batches 1, 2, 4 and 8 is reported under Memory.

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-CUDA-TF32",
    revision="7fb4118b83a07f3af901ee9d52f3f3af67224c35",
    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-CUDA-TF32",
    revision="7d2e54941a3c2c9d1ac4af27f21531bf39dfb1dd",
    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.