AutomatosX commited on
Commit
1b4ed07
·
verified ·
1 Parent(s): 7efd000

Publish pinned CUDA TF32 inference source and example

Browse files
inference/LICENSE ADDED
@@ -0,0 +1,661 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ GNU AFFERO GENERAL PUBLIC LICENSE
2
+ Version 3, 19 November 2007
3
+
4
+ Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
5
+ Everyone is permitted to copy and distribute verbatim copies
6
+ of this license document, but changing it is not allowed.
7
+
8
+ Preamble
9
+
10
+ The GNU Affero General Public License is a free, copyleft license for
11
+ software and other kinds of works, specifically designed to ensure
12
+ cooperation with the community in the case of network server software.
13
+
14
+ The licenses for most software and other practical works are designed
15
+ to take away your freedom to share and change the works. By contrast,
16
+ our General Public Licenses are intended to guarantee your freedom to
17
+ share and change all versions of a program--to make sure it remains free
18
+ software for all its users.
19
+
20
+ When we speak of free software, we are referring to freedom, not
21
+ price. Our General Public Licenses are designed to make sure that you
22
+ have the freedom to distribute copies of free software (and charge for
23
+ them if you wish), that you receive source code or can get it if you
24
+ want it, that you can change the software or use pieces of it in new
25
+ free programs, and that you know you can do these things.
26
+
27
+ Developers that use our General Public Licenses protect your rights
28
+ with two steps: (1) assert copyright on the software, and (2) offer
29
+ you this License which gives you legal permission to copy, distribute
30
+ and/or modify the software.
31
+
32
+ A secondary benefit of defending all users' freedom is that
33
+ improvements made in alternate versions of the program, if they
34
+ receive widespread use, become available for other developers to
35
+ incorporate. Many developers of free software are heartened and
36
+ encouraged by the resulting cooperation. However, in the case of
37
+ software used on network servers, this result may fail to come about.
38
+ The GNU General Public License permits making a modified version and
39
+ letting the public access it on a server without ever releasing its
40
+ source code to the public.
41
+
42
+ The GNU Affero General Public License is designed specifically to
43
+ ensure that, in such cases, the modified source code becomes available
44
+ to the community. It requires the operator of a network server to
45
+ provide the source code of the modified version running there to the
46
+ users of that server. Therefore, public use of a modified version, on
47
+ a publicly accessible server, gives the public access to the source
48
+ code of the modified version.
49
+
50
+ An older license, called the Affero General Public License and
51
+ published by Affero, was designed to accomplish similar goals. This is
52
+ a different license, not a version of the Affero GPL, but Affero has
53
+ released a new version of the Affero GPL which permits relicensing under
54
+ this license.
55
+
56
+ The precise terms and conditions for copying, distribution and
57
+ modification follow.
58
+
59
+ TERMS AND CONDITIONS
60
+
61
+ 0. Definitions.
62
+
63
+ "This License" refers to version 3 of the GNU Affero General Public License.
64
+
65
+ "Copyright" also means copyright-like laws that apply to other kinds of
66
+ works, such as semiconductor masks.
67
+
68
+ "The Program" refers to any copyrightable work licensed under this
69
+ License. Each licensee is addressed as "you". "Licensees" and
70
+ "recipients" may be individuals or organizations.
71
+
72
+ To "modify" a work means to copy from or adapt all or part of the work
73
+ in a fashion requiring copyright permission, other than the making of an
74
+ exact copy. The resulting work is called a "modified version" of the
75
+ earlier work or a work "based on" the earlier work.
76
+
77
+ A "covered work" means either the unmodified Program or a work based
78
+ on the Program.
79
+
80
+ To "propagate" a work means to do anything with it that, without
81
+ permission, would make you directly or secondarily liable for
82
+ infringement under applicable copyright law, except executing it on a
83
+ computer or modifying a private copy. Propagation includes copying,
84
+ distribution (with or without modification), making available to the
85
+ public, and in some countries other activities as well.
86
+
87
+ To "convey" a work means any kind of propagation that enables other
88
+ parties to make or receive copies. Mere interaction with a user through
89
+ a computer network, with no transfer of a copy, is not conveying.
90
+
91
+ An interactive user interface displays "Appropriate Legal Notices"
92
+ to the extent that it includes a convenient and prominently visible
93
+ feature that (1) displays an appropriate copyright notice, and (2)
94
+ tells the user that there is no warranty for the work (except to the
95
+ extent that warranties are provided), that licensees may convey the
96
+ work under this License, and how to view a copy of this License. If
97
+ the interface presents a list of user commands or options, such as a
98
+ menu, a prominent item in the list meets this criterion.
99
+
100
+ 1. Source Code.
101
+
102
+ The "source code" for a work means the preferred form of the work
103
+ for making modifications to it. "Object code" means any non-source
104
+ form of a work.
105
+
106
+ A "Standard Interface" means an interface that either is an official
107
+ standard defined by a recognized standards body, or, in the case of
108
+ interfaces specified for a particular programming language, one that
109
+ is widely used among developers working in that language.
110
+
111
+ The "System Libraries" of an executable work include anything, other
112
+ than the work as a whole, that (a) is included in the normal form of
113
+ packaging a Major Component, but which is not part of that Major
114
+ Component, and (b) serves only to enable use of the work with that
115
+ Major Component, or to implement a Standard Interface for which an
116
+ implementation is available to the public in source code form. A
117
+ "Major Component", in this context, means a major essential component
118
+ (kernel, window system, and so on) of the specific operating system
119
+ (if any) on which the executable work runs, or a compiler used to
120
+ produce the work, or an object code interpreter used to run it.
121
+
122
+ The "Corresponding Source" for a work in object code form means all
123
+ the source code needed to generate, install, and (for an executable
124
+ work) run the object code and to modify the work, including scripts to
125
+ control those activities. However, it does not include the work's
126
+ System Libraries, or general-purpose tools or generally available free
127
+ programs which are used unmodified in performing those activities but
128
+ which are not part of the work. For example, Corresponding Source
129
+ includes interface definition files associated with source files for
130
+ the work, and the source code for shared libraries and dynamically
131
+ linked subprograms that the work is specifically designed to require,
132
+ such as by intimate data communication or control flow between those
133
+ subprograms and other parts of the work.
134
+
135
+ The Corresponding Source need not include anything that users
136
+ can regenerate automatically from other parts of the Corresponding
137
+ Source.
138
+
139
+ The Corresponding Source for a work in source code form is that
140
+ same work.
141
+
142
+ 2. Basic Permissions.
143
+
144
+ All rights granted under this License are granted for the term of
145
+ copyright on the Program, and are irrevocable provided the stated
146
+ conditions are met. This License explicitly affirms your unlimited
147
+ permission to run the unmodified Program. The output from running a
148
+ covered work is covered by this License only if the output, given its
149
+ content, constitutes a covered work. This License acknowledges your
150
+ rights of fair use or other equivalent, as provided by copyright law.
151
+
152
+ You may make, run and propagate covered works that you do not
153
+ convey, without conditions so long as your license otherwise remains
154
+ in force. You may convey covered works to others for the sole purpose
155
+ of having them make modifications exclusively for you, or provide you
156
+ with facilities for running those works, provided that you comply with
157
+ the terms of this License in conveying all material for which you do
158
+ not control copyright. Those thus making or running the covered works
159
+ for you must do so exclusively on your behalf, under your direction
160
+ and control, on terms that prohibit them from making any copies of
161
+ your copyrighted material outside their relationship with you.
162
+
163
+ Conveying under any other circumstances is permitted solely under
164
+ the conditions stated below. Sublicensing is not allowed; section 10
165
+ makes it unnecessary.
166
+
167
+ 3. Protecting Users' Legal Rights From Anti-Circumvention Law.
168
+
169
+ No covered work shall be deemed part of an effective technological
170
+ measure under any applicable law fulfilling obligations under article
171
+ 11 of the WIPO copyright treaty adopted on 20 December 1996, or
172
+ similar laws prohibiting or restricting circumvention of such
173
+ measures.
174
+
175
+ When you convey a covered work, you waive any legal power to forbid
176
+ circumvention of technological measures to the extent such circumvention
177
+ is effected by exercising rights under this License with respect to
178
+ the covered work, and you disclaim any intention to limit operation or
179
+ modification of the work as a means of enforcing, against the work's
180
+ users, your or third parties' legal rights to forbid circumvention of
181
+ technological measures.
182
+
183
+ 4. Conveying Verbatim Copies.
184
+
185
+ You may convey verbatim copies of the Program's source code as you
186
+ receive it, in any medium, provided that you conspicuously and
187
+ appropriately publish on each copy an appropriate copyright notice;
188
+ keep intact all notices stating that this License and any
189
+ non-permissive terms added in accord with section 7 apply to the code;
190
+ keep intact all notices of the absence of any warranty; and give all
191
+ recipients a copy of this License along with the Program.
192
+
193
+ You may charge any price or no price for each copy that you convey,
194
+ and you may offer support or warranty protection for a fee.
195
+
196
+ 5. Conveying Modified Source Versions.
197
+
198
+ You may convey a work based on the Program, or the modifications to
199
+ produce it from the Program, in the form of source code under the
200
+ terms of section 4, provided that you also meet all of these conditions:
201
+
202
+ a) The work must carry prominent notices stating that you modified
203
+ it, and giving a relevant date.
204
+
205
+ b) The work must carry prominent notices stating that it is
206
+ released under this License and any conditions added under section
207
+ 7. This requirement modifies the requirement in section 4 to
208
+ "keep intact all notices".
209
+
210
+ c) You must license the entire work, as a whole, under this
211
+ License to anyone who comes into possession of a copy. This
212
+ License will therefore apply, along with any applicable section 7
213
+ additional terms, to the whole of the work, and all its parts,
214
+ regardless of how they are packaged. This License gives no
215
+ permission to license the work in any other way, but it does not
216
+ invalidate such permission if you have separately received it.
217
+
218
+ d) If the work has interactive user interfaces, each must display
219
+ Appropriate Legal Notices; however, if the Program has interactive
220
+ interfaces that do not display Appropriate Legal Notices, your
221
+ work need not make them do so.
222
+
223
+ A compilation of a covered work with other separate and independent
224
+ works, which are not by their nature extensions of the covered work,
225
+ and which are not combined with it such as to form a larger program,
226
+ in or on a volume of a storage or distribution medium, is called an
227
+ "aggregate" if the compilation and its resulting copyright are not
228
+ used to limit the access or legal rights of the compilation's users
229
+ beyond what the individual works permit. Inclusion of a covered work
230
+ in an aggregate does not cause this License to apply to the other
231
+ parts of the aggregate.
232
+
233
+ 6. Conveying Non-Source Forms.
234
+
235
+ You may convey a covered work in object code form under the terms
236
+ of sections 4 and 5, provided that you also convey the
237
+ machine-readable Corresponding Source under the terms of this License,
238
+ in one of these ways:
239
+
240
+ a) Convey the object code in, or embodied in, a physical product
241
+ (including a physical distribution medium), accompanied by the
242
+ Corresponding Source fixed on a durable physical medium
243
+ customarily used for software interchange.
244
+
245
+ b) Convey the object code in, or embodied in, a physical product
246
+ (including a physical distribution medium), accompanied by a
247
+ written offer, valid for at least three years and valid for as
248
+ long as you offer spare parts or customer support for that product
249
+ model, to give anyone who possesses the object code either (1) a
250
+ copy of the Corresponding Source for all the software in the
251
+ product that is covered by this License, on a durable physical
252
+ medium customarily used for software interchange, for a price no
253
+ more than your reasonable cost of physically performing this
254
+ conveying of source, or (2) access to copy the
255
+ Corresponding Source from a network server at no charge.
256
+
257
+ c) Convey individual copies of the object code with a copy of the
258
+ written offer to provide the Corresponding Source. This
259
+ alternative is allowed only occasionally and noncommercially, and
260
+ only if you received the object code with such an offer, in accord
261
+ with subsection 6b.
262
+
263
+ d) Convey the object code by offering access from a designated
264
+ place (gratis or for a charge), and offer equivalent access to the
265
+ Corresponding Source in the same way through the same place at no
266
+ further charge. You need not require recipients to copy the
267
+ Corresponding Source along with the object code. If the place to
268
+ copy the object code is a network server, the Corresponding Source
269
+ may be on a different server (operated by you or a third party)
270
+ that supports equivalent copying facilities, provided you maintain
271
+ clear directions next to the object code saying where to find the
272
+ Corresponding Source. Regardless of what server hosts the
273
+ Corresponding Source, you remain obligated to ensure that it is
274
+ available for as long as needed to satisfy these requirements.
275
+
276
+ e) Convey the object code using peer-to-peer transmission, provided
277
+ you inform other peers where the object code and Corresponding
278
+ Source of the work are being offered to the general public at no
279
+ charge under subsection 6d.
280
+
281
+ A separable portion of the object code, whose source code is excluded
282
+ from the Corresponding Source as a System Library, need not be
283
+ included in conveying the object code work.
284
+
285
+ A "User Product" is either (1) a "consumer product", which means any
286
+ tangible personal property which is normally used for personal, family,
287
+ or household purposes, or (2) anything designed or sold for incorporation
288
+ into a dwelling. In determining whether a product is a consumer product,
289
+ doubtful cases shall be resolved in favor of coverage. For a particular
290
+ product received by a particular user, "normally used" refers to a
291
+ typical or common use of that class of product, regardless of the status
292
+ of the particular user or of the way in which the particular user
293
+ actually uses, or expects or is expected to use, the product. A product
294
+ is a consumer product regardless of whether the product has substantial
295
+ commercial, industrial or non-consumer uses, unless such uses represent
296
+ the only significant mode of use of the product.
297
+
298
+ "Installation Information" for a User Product means any methods,
299
+ procedures, authorization keys, or other information required to install
300
+ and execute modified versions of a covered work in that User Product from
301
+ a modified version of its Corresponding Source. The information must
302
+ suffice to ensure that the continued functioning of the modified object
303
+ code is in no case prevented or interfered with solely because
304
+ modification has been made.
305
+
306
+ If you convey an object code work under this section in, or with, or
307
+ specifically for use in, a User Product, and the conveying occurs as
308
+ part of a transaction in which the right of possession and use of the
309
+ User Product is transferred to the recipient in perpetuity or for a
310
+ fixed term (regardless of how the transaction is characterized), the
311
+ Corresponding Source conveyed under this section must be accompanied
312
+ by the Installation Information. But this requirement does not apply
313
+ if neither you nor any third party retains the ability to install
314
+ modified object code on the User Product (for example, the work has
315
+ been installed in ROM).
316
+
317
+ The requirement to provide Installation Information does not include a
318
+ requirement to continue to provide support service, warranty, or updates
319
+ for a work that has been modified or installed by the recipient, or for
320
+ the User Product in which it has been modified or installed. Access to a
321
+ network may be denied when the modification itself materially and
322
+ adversely affects the operation of the network or violates the rules and
323
+ protocols for communication across the network.
324
+
325
+ Corresponding Source conveyed, and Installation Information provided,
326
+ in accord with this section must be in a format that is publicly
327
+ documented (and with an implementation available to the public in
328
+ source code form), and must require no special password or key for
329
+ unpacking, reading or copying.
330
+
331
+ 7. Additional Terms.
332
+
333
+ "Additional permissions" are terms that supplement the terms of this
334
+ License by making exceptions from one or more of its conditions.
335
+ Additional permissions that are applicable to the entire Program shall
336
+ be treated as though they were included in this License, to the extent
337
+ that they are valid under applicable law. If additional permissions
338
+ apply only to part of the Program, that part may be used separately
339
+ under those permissions, but the entire Program remains governed by
340
+ this License without regard to the additional permissions.
341
+
342
+ When you convey a copy of a covered work, you may at your option
343
+ remove any additional permissions from that copy, or from any part of
344
+ it. (Additional permissions may be written to require their own
345
+ removal in certain cases when you modify the work.) You may place
346
+ additional permissions on material, added by you to a covered work,
347
+ for which you have or can give appropriate copyright permission.
348
+
349
+ Notwithstanding any other provision of this License, for material you
350
+ add to a covered work, you may (if authorized by the copyright holders of
351
+ that material) supplement the terms of this License with terms:
352
+
353
+ a) Disclaiming warranty or limiting liability differently from the
354
+ terms of sections 15 and 16 of this License; or
355
+
356
+ b) Requiring preservation of specified reasonable legal notices or
357
+ author attributions in that material or in the Appropriate Legal
358
+ Notices displayed by works containing it; or
359
+
360
+ c) Prohibiting misrepresentation of the origin of that material, or
361
+ requiring that modified versions of such material be marked in
362
+ reasonable ways as different from the original version; or
363
+
364
+ d) Limiting the use for publicity purposes of names of licensors or
365
+ authors of the material; or
366
+
367
+ e) Declining to grant rights under trademark law for use of some
368
+ trade names, trademarks, or service marks; or
369
+
370
+ f) Requiring indemnification of licensors and authors of that
371
+ material by anyone who conveys the material (or modified versions of
372
+ it) with contractual assumptions of liability to the recipient, for
373
+ any liability that these contractual assumptions directly impose on
374
+ those licensors and authors.
375
+
376
+ All other non-permissive additional terms are considered "further
377
+ restrictions" within the meaning of section 10. If the Program as you
378
+ received it, or any part of it, contains a notice stating that it is
379
+ governed by this License along with a term that is a further
380
+ restriction, you may remove that term. If a license document contains
381
+ a further restriction but permits relicensing or conveying under this
382
+ License, you may add to a covered work material governed by the terms
383
+ of that license document, provided that the further restriction does
384
+ not survive such relicensing or conveying.
385
+
386
+ If you add terms to a covered work in accord with this section, you
387
+ must place, in the relevant source files, a statement of the
388
+ additional terms that apply to those files, or a notice indicating
389
+ where to find the applicable terms.
390
+
391
+ Additional terms, permissive or non-permissive, may be stated in the
392
+ form of a separately written license, or stated as exceptions;
393
+ the above requirements apply either way.
394
+
395
+ 8. Termination.
396
+
397
+ You may not propagate or modify a covered work except as expressly
398
+ provided under this License. Any attempt otherwise to propagate or
399
+ modify it is void, and will automatically terminate your rights under
400
+ this License (including any patent licenses granted under the third
401
+ paragraph of section 11).
402
+
403
+ However, if you cease all violation of this License, then your
404
+ license from a particular copyright holder is reinstated (a)
405
+ provisionally, unless and until the copyright holder explicitly and
406
+ finally terminates your license, and (b) permanently, if the copyright
407
+ holder fails to notify you of the violation by some reasonable means
408
+ prior to 60 days after the cessation.
409
+
410
+ Moreover, your license from a particular copyright holder is
411
+ reinstated permanently if the copyright holder notifies you of the
412
+ violation by some reasonable means, this is the first time you have
413
+ received notice of violation of this License (for any work) from that
414
+ copyright holder, and you cure the violation prior to 30 days after
415
+ your receipt of the notice.
416
+
417
+ Termination of your rights under this section does not terminate the
418
+ licenses of parties who have received copies or rights from you under
419
+ this License. If your rights have been terminated and not permanently
420
+ reinstated, you do not qualify to receive new licenses for the same
421
+ material under section 10.
422
+
423
+ 9. Acceptance Not Required for Having Copies.
424
+
425
+ You are not required to accept this License in order to receive or
426
+ run a copy of the Program. Ancillary propagation of a covered work
427
+ occurring solely as a consequence of using peer-to-peer transmission
428
+ to receive a copy likewise does not require acceptance. However,
429
+ nothing other than this License grants you permission to propagate or
430
+ modify any covered work. These actions infringe copyright if you do
431
+ not accept this License. Therefore, by modifying or propagating a
432
+ covered work, you indicate your acceptance of this License to do so.
433
+
434
+ 10. Automatic Licensing of Downstream Recipients.
435
+
436
+ Each time you convey a covered work, the recipient automatically
437
+ receives a license from the original licensors, to run, modify and
438
+ propagate that work, subject to this License. You are not responsible
439
+ for enforcing compliance by third parties with this License.
440
+
441
+ An "entity transaction" is a transaction transferring control of an
442
+ organization, or substantially all assets of one, or subdividing an
443
+ organization, or merging organizations. If propagation of a covered
444
+ work results from an entity transaction, each party to that
445
+ transaction who receives a copy of the work also receives whatever
446
+ licenses to the work the party's predecessor in interest had or could
447
+ give under the previous paragraph, plus a right to possession of the
448
+ Corresponding Source of the work from the predecessor in interest, if
449
+ the predecessor has it or can get it with reasonable efforts.
450
+
451
+ You may not impose any further restrictions on the exercise of the
452
+ rights granted or affirmed under this License. For example, you may
453
+ not impose a license fee, royalty, or other charge for exercise of
454
+ rights granted under this License, and you may not initiate litigation
455
+ (including a cross-claim or counterclaim in a lawsuit) alleging that
456
+ any patent claim is infringed by making, using, selling, offering for
457
+ sale, or importing the Program or any portion of it.
458
+
459
+ 11. Patents.
460
+
461
+ A "contributor" is a copyright holder who authorizes use under this
462
+ License of the Program or a work on which the Program is based. The
463
+ work thus licensed is called the contributor's "contributor version".
464
+
465
+ A contributor's "essential patent claims" are all patent claims
466
+ owned or controlled by the contributor, whether already acquired or
467
+ hereafter acquired, that would be infringed by some manner, permitted
468
+ by this License, of making, using, or selling its contributor version,
469
+ but do not include claims that would be infringed only as a
470
+ consequence of further modification of the contributor version. For
471
+ purposes of this definition, "control" includes the right to grant
472
+ patent sublicenses in a manner consistent with the requirements of
473
+ this License.
474
+
475
+ Each contributor grants you a non-exclusive, worldwide, royalty-free
476
+ patent license under the contributor's essential patent claims, to
477
+ make, use, sell, offer for sale, import and otherwise run, modify and
478
+ propagate the contents of its contributor version.
479
+
480
+ In the following three paragraphs, a "patent license" is any express
481
+ agreement or commitment, however denominated, not to enforce a patent
482
+ (such as an express permission to practice a patent or covenant not to
483
+ sue for patent infringement). To "grant" such a patent license to a
484
+ party means to make such an agreement or commitment not to enforce a
485
+ patent against the party.
486
+
487
+ If you convey a covered work, knowingly relying on a patent license,
488
+ and the Corresponding Source of the work is not available for anyone
489
+ to copy, free of charge and under the terms of this License, through a
490
+ publicly available network server or other readily accessible means,
491
+ then you must either (1) cause the Corresponding Source to be so
492
+ available, or (2) arrange to deprive yourself of the benefit of the
493
+ patent license for this particular work, or (3) arrange, in a manner
494
+ consistent with the requirements of this License, to extend the patent
495
+ license to downstream recipients. "Knowingly relying" means you have
496
+ actual knowledge that, but for the patent license, your conveying the
497
+ covered work in a country, or your recipient's use of the covered work
498
+ in a country, would infringe one or more identifiable patents in that
499
+ country that you have reason to believe are valid.
500
+
501
+ If, pursuant to or in connection with a single transaction or
502
+ arrangement, you convey, or propagate by procuring conveyance of, a
503
+ covered work, and grant a patent license to some of the parties
504
+ receiving the covered work authorizing them to use, propagate, modify
505
+ or convey a specific copy of the covered work, then the patent license
506
+ you grant is automatically extended to all recipients of the covered
507
+ work and works based on it.
508
+
509
+ A patent license is "discriminatory" if it does not include within
510
+ the scope of its coverage, prohibits the exercise of, or is
511
+ conditioned on the non-exercise of one or more of the rights that are
512
+ specifically granted under this License. You may not convey a covered
513
+ work if you are a party to an arrangement with a third party that is
514
+ in the business of distributing software, under which you make payment
515
+ to the third party based on the extent of your activity of conveying
516
+ the work, and under which the third party grants, to any of the
517
+ parties who would receive the covered work from you, a discriminatory
518
+ patent license (a) in connection with copies of the covered work
519
+ conveyed by you (or copies made from those copies), or (b) primarily
520
+ for and in connection with specific products or compilations that
521
+ contain the covered work, unless you entered into that arrangement,
522
+ or that patent license was granted, prior to 28 March 2007.
523
+
524
+ Nothing in this License shall be construed as excluding or limiting
525
+ any implied license or other defenses to infringement that may
526
+ otherwise be available to you under applicable patent law.
527
+
528
+ 12. No Surrender of Others' Freedom.
529
+
530
+ If conditions are imposed on you (whether by court order, agreement or
531
+ otherwise) that contradict the conditions of this License, they do not
532
+ excuse you from the conditions of this License. If you cannot convey a
533
+ covered work so as to satisfy simultaneously your obligations under this
534
+ License and any other pertinent obligations, then as a consequence you may
535
+ not convey it at all. For example, if you agree to terms that obligate you
536
+ to collect a royalty for further conveying from those to whom you convey
537
+ the Program, the only way you could satisfy both those terms and this
538
+ License would be to refrain entirely from conveying the Program.
539
+
540
+ 13. Remote Network Interaction; Use with the GNU General Public License.
541
+
542
+ Notwithstanding any other provision of this License, if you modify the
543
+ Program, your modified version must prominently offer all users
544
+ interacting with it remotely through a computer network (if your version
545
+ supports such interaction) an opportunity to receive the Corresponding
546
+ Source of your version by providing access to the Corresponding Source
547
+ from a network server at no charge, through some standard or customary
548
+ means of facilitating copying of software. This Corresponding Source
549
+ shall include the Corresponding Source for any work covered by version 3
550
+ of the GNU General Public License that is incorporated pursuant to the
551
+ following paragraph.
552
+
553
+ Notwithstanding any other provision of this License, you have
554
+ permission to link or combine any covered work with a work licensed
555
+ under version 3 of the GNU General Public License into a single
556
+ combined work, and to convey the resulting work. The terms of this
557
+ License will continue to apply to the part which is the covered work,
558
+ but the work with which it is combined will remain governed by version
559
+ 3 of the GNU General Public License.
560
+
561
+ 14. Revised Versions of this License.
562
+
563
+ The Free Software Foundation may publish revised and/or new versions of
564
+ the GNU Affero General Public License from time to time. Such new versions
565
+ will be similar in spirit to the present version, but may differ in detail to
566
+ address new problems or concerns.
567
+
568
+ Each version is given a distinguishing version number. If the
569
+ Program specifies that a certain numbered version of the GNU Affero General
570
+ Public License "or any later version" applies to it, you have the
571
+ option of following the terms and conditions either of that numbered
572
+ version or of any later version published by the Free Software
573
+ Foundation. If the Program does not specify a version number of the
574
+ GNU Affero General Public License, you may choose any version ever published
575
+ by the Free Software Foundation.
576
+
577
+ If the Program specifies that a proxy can decide which future
578
+ versions of the GNU Affero General Public License can be used, that proxy's
579
+ public statement of acceptance of a version permanently authorizes you
580
+ to choose that version for the Program.
581
+
582
+ Later license versions may give you additional or different
583
+ permissions. However, no additional obligations are imposed on any
584
+ author or copyright holder as a result of your choosing to follow a
585
+ later version.
586
+
587
+ 15. Disclaimer of Warranty.
588
+
589
+ THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
590
+ APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
591
+ HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
592
+ OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
593
+ THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
594
+ PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
595
+ IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
596
+ ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
597
+
598
+ 16. Limitation of Liability.
599
+
600
+ IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
601
+ WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
602
+ THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
603
+ GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
604
+ USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
605
+ DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
606
+ PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
607
+ EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
608
+ SUCH DAMAGES.
609
+
610
+ 17. Interpretation of Sections 15 and 16.
611
+
612
+ If the disclaimer of warranty and limitation of liability provided
613
+ above cannot be given local legal effect according to their terms,
614
+ reviewing courts shall apply local law that most closely approximates
615
+ an absolute waiver of all civil liability in connection with the
616
+ Program, unless a warranty or assumption of liability accompanies a
617
+ copy of the Program in return for a fee.
618
+
619
+ END OF TERMS AND CONDITIONS
620
+
621
+ How to Apply These Terms to Your New Programs
622
+
623
+ If you develop a new program, and you want it to be of the greatest
624
+ possible use to the public, the best way to achieve this is to make it
625
+ free software which everyone can redistribute and change under these terms.
626
+
627
+ To do so, attach the following notices to the program. It is safest
628
+ to attach them to the start of each source file to most effectively
629
+ state the exclusion of warranty; and each file should have at least
630
+ the "copyright" line and a pointer to where the full notice is found.
631
+
632
+ <one line to give the program's name and a brief idea of what it does.>
633
+ Copyright (C) <year> <name of author>
634
+
635
+ This program is free software: you can redistribute it and/or modify
636
+ it under the terms of the GNU Affero General Public License as published by
637
+ the Free Software Foundation, either version 3 of the License, or
638
+ (at your option) any later version.
639
+
640
+ This program is distributed in the hope that it will be useful,
641
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
642
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
643
+ GNU Affero General Public License for more details.
644
+
645
+ You should have received a copy of the GNU Affero General Public License
646
+ along with this program. If not, see <https://www.gnu.org/licenses/>.
647
+
648
+ Also add information on how to contact you by electronic and paper mail.
649
+
650
+ If your software can interact with users remotely through a computer
651
+ network, you should also make sure that it provides a way for users to
652
+ get its source. For example, if your program is a web application, its
653
+ interface could display a "Source" link that leads users to an archive
654
+ of the code. There are many ways you could offer source, and different
655
+ solutions will be better for different programs; see section 13 for the
656
+ specific requirements.
657
+
658
+ You should also get your employer (if you work as a programmer) or school,
659
+ if any, to sign a "copyright disclaimer" for the program, if necessary.
660
+ For more information on this, and how to apply and follow the GNU AGPL, see
661
+ <https://www.gnu.org/licenses/>.
inference/NOTICE ADDED
@@ -0,0 +1,246 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ AXDocLayout-M v1.2 CUDA TF32 runtime update
2
+ Modified by AutomatosX on 2026-10-17: explicit ONNX Runtime CUDA execution,
3
+ TF32 policy and placement verification, equivalent vectorized postprocessing,
4
+ release identity, integrity bindings and inference documentation. Model weights,
5
+ public classes and raster coordinates are unchanged from v1.1. The complete
6
+ inference source is supplied in inference/. License: AGPL-3.0-only.
7
+
8
+ AXDocLayout-M v1.2
9
+ Copyright 2026 AutomatosX
10
+
11
+ This release as a whole, including its model weights and modified inference
12
+ source, is licensed under GNU AGPL version 3 only (AGPL-3.0-only).
13
+ See LICENSE. Commercial use is permitted subject to the license terms.
14
+ No additional non-commercial restriction is imposed.
15
+
16
+ Modified by AutomatosX on 2026-10-17: updated model artifacts, release
17
+ identity, source-neutral taxonomy identifier, artifact integrity bindings,
18
+ usage documentation and distribution license. Public class IDs and raster
19
+ coordinates are unchanged. The complete inference source is in inference/.
20
+
21
+ Previously Apache-2.0-licensed components retain their original notices and
22
+ rights. The following notice and full Apache-2.0 text are preserved for those
23
+ components; they do not offer the new model weights under Apache-2.0.
24
+
25
+ --- Preserved component notice ---
26
+ AX DocLayout
27
+ Copyright 2026 AutomatosX
28
+
29
+ Licensed under the Apache License, Version 2.0 (the "License");
30
+ you may not use this file except in compliance with the License.
31
+ You may obtain a copy of the License at
32
+
33
+ http://www.apache.org/licenses/LICENSE-2.0
34
+
35
+ Unless required by applicable law or agreed to in writing, software
36
+ distributed under the License is distributed on an "AS IS" BASIS,
37
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
38
+ See the License for the specific language governing permissions and
39
+ limitations under the License.
40
+
41
+ The detector architecture is RT-DETRv2 (`rtdetrv2_r50vd_m`). Upstream
42
+ RT-DETR is licensed under Apache-2.0 and is not included in this bundle.
43
+
44
+ --- Apache License 2.0 for prior components ---
45
+
46
+ Apache License
47
+ Version 2.0, January 2004
48
+ http://www.apache.org/licenses/
49
+
50
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
51
+
52
+ 1. Definitions.
53
+
54
+ "License" shall mean the terms and conditions for use, reproduction,
55
+ and distribution as defined by Sections 1 through 9 of this document.
56
+
57
+ "Licensor" shall mean the copyright owner or entity authorized by
58
+ the copyright owner that is granting the License.
59
+
60
+ "Legal Entity" shall mean the union of the acting entity and all
61
+ other entities that control, are controlled by, or are under common
62
+ control with that entity. For the purposes of this definition,
63
+ "control" means (i) the power, direct or indirect, to cause the
64
+ direction or management of such entity, whether by contract or
65
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
66
+ outstanding shares, or (iii) beneficial ownership of such entity.
67
+
68
+ "You" (or "Your") shall mean an individual or Legal Entity
69
+ exercising permissions granted by this License.
70
+
71
+ "Source" form shall mean the preferred form for making modifications,
72
+ including but not limited to software source code, documentation
73
+ source, and configuration files.
74
+
75
+ "Object" form shall mean any form resulting from mechanical
76
+ transformation or translation of a Source form, including but
77
+ not limited to compiled object code, generated documentation,
78
+ and conversions to other media types.
79
+
80
+ "Work" shall mean the work of authorship, whether in Source or
81
+ Object form, made available under the License, as indicated by a
82
+ copyright notice that is included in or attached to the work
83
+ (an example is provided in the Appendix below).
84
+
85
+ "Derivative Works" shall mean any work, whether in Source or Object
86
+ form, that is based on (or derived from) the Work and for which the
87
+ editorial revisions, annotations, elaborations, or other modifications
88
+ represent, as a whole, an original work of authorship. For the purposes
89
+ of this License, Derivative Works shall not include works that remain
90
+ separable from, or merely link (or bind by name) to the interfaces of,
91
+ the Work and Derivative Works thereof.
92
+
93
+ "Contribution" shall mean any work of authorship, including
94
+ the original version of the Work and any modifications or additions
95
+ to that Work or Derivative Works thereof, that is intentionally
96
+ submitted to Licensor for inclusion in the Work by the copyright owner
97
+ or by an individual or Legal Entity authorized to submit on behalf of
98
+ the copyright owner. For the purposes of this definition, "submitted"
99
+ means any form of electronic, verbal, or written communication sent
100
+ to the Licensor or its representatives, including but not limited to
101
+ communication on electronic mailing lists, source code control systems,
102
+ and issue tracking systems that are managed by, or on behalf of, the
103
+ Licensor for the purpose of discussing and improving the Work, but
104
+ excluding communication that is conspicuously marked or otherwise
105
+ designated in writing by the copyright owner as "Not a Contribution."
106
+
107
+ "Contributor" shall mean Licensor and any individual or Legal Entity
108
+ on behalf of whom a Contribution has been received by Licensor and
109
+ subsequently incorporated within the Work.
110
+
111
+ 2. Grant of Copyright License. Subject to the terms and conditions of
112
+ this License, each Contributor hereby grants to You a perpetual,
113
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
114
+ copyright license to reproduce, prepare Derivative Works of,
115
+ publicly display, publicly perform, sublicense, and distribute the
116
+ Work and such Derivative Works in Source or Object form.
117
+
118
+ 3. Grant of Patent License. Subject to the terms and conditions of
119
+ this License, each Contributor hereby grants to You a perpetual,
120
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
121
+ (except as stated in this section) patent license to make, have made,
122
+ use, offer to sell, sell, import, and otherwise transfer the Work,
123
+ where such license applies only to those patent claims licensable
124
+ by such Contributor that are necessarily infringed by their
125
+ Contribution(s) alone or by combination of their Contribution(s)
126
+ with the Work to which such Contribution(s) was submitted. If You
127
+ institute patent litigation against any entity (including a
128
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
129
+ or a Contribution incorporated within the Work constitutes direct
130
+ or contributory patent infringement, then any patent licenses
131
+ granted to You under this License for that Work shall terminate
132
+ as of the date such litigation is filed.
133
+
134
+ 4. Redistribution. You may reproduce and distribute copies of the
135
+ Work or Derivative Works thereof in any medium, with or without
136
+ modifications, and in Source or Object form, provided that You
137
+ meet the following conditions:
138
+
139
+ (a) You must give any other recipients of the Work or
140
+ Derivative Works a copy of this License; and
141
+
142
+ (b) You must cause any modified files to carry prominent notices
143
+ stating that You changed the files; and
144
+
145
+ (c) You must retain, in the Source form of any Derivative Works
146
+ that You distribute, all copyright, patent, trademark, and
147
+ attribution notices from the Source form of the Work,
148
+ excluding those notices that do not pertain to any part of
149
+ the Derivative Works; and
150
+
151
+ (d) If the Work includes a "NOTICE" text file as part of its
152
+ distribution, then any Derivative Works that You distribute must
153
+ include a readable copy of the attribution notices contained
154
+ within such NOTICE file, excluding those notices that do not
155
+ pertain to any part of the Derivative Works, in at least one
156
+ of the following places: within a NOTICE text file distributed
157
+ as part of the Derivative Works; within the Source form or
158
+ documentation, if provided along with the Derivative Works; or,
159
+ within a display generated by the Derivative Works, if and
160
+ wherever such third-party notices normally appear. The contents
161
+ of the NOTICE file are for informational purposes only and
162
+ do not modify the License. You may add Your own attribution
163
+ notices within Derivative Works that You distribute, alongside
164
+ or as an addendum to the NOTICE text from the Work, provided
165
+ that such additional attribution notices cannot be construed
166
+ as modifying the License.
167
+
168
+ You may add Your own copyright statement to Your modifications and
169
+ may provide additional or different license terms and conditions
170
+ for use, reproduction, or distribution of Your modifications, or
171
+ for any such Derivative Works as a whole, provided Your use,
172
+ reproduction, and distribution of the Work otherwise complies with
173
+ the conditions stated in this License.
174
+
175
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
176
+ any Contribution intentionally submitted for inclusion in the Work
177
+ by You to the Licensor shall be under the terms and conditions of
178
+ this License, without any additional terms or conditions.
179
+ Notwithstanding the above, nothing herein shall supersede or modify
180
+ the terms of any separate license agreement you may have executed
181
+ with Licensor regarding such Contributions.
182
+
183
+ 6. Trademarks. This License does not grant permission to use the trade
184
+ names, trademarks, service marks, or product names of the Licensor,
185
+ except as required for reasonable and customary use in describing the
186
+ origin of the Work and reproducing the content of the NOTICE file.
187
+
188
+ 7. Disclaimer of Warranty. Unless required by applicable law or
189
+ agreed to in writing, Licensor provides the Work (and each
190
+ Contributor provides its Contributions) on an "AS IS" BASIS,
191
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
192
+ implied, including, without limitation, any warranties or conditions
193
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
194
+ PARTICULAR PURPOSE. You are solely responsible for determining the
195
+ appropriateness of using or redistributing the Work and assume any
196
+ risks associated with Your exercise of permissions under this License.
197
+
198
+ 8. Limitation of Liability. In no event and under no legal theory,
199
+ whether in tort (including negligence), contract, or otherwise,
200
+ unless required by applicable law (such as deliberate and grossly
201
+ negligent acts) or agreed to in writing, shall any Contributor be
202
+ liable to You for damages, including any direct, indirect, special,
203
+ incidental, or consequential damages of any character arising as a
204
+ result of this License or out of the use or inability to use the
205
+ Work (including but not limited to damages for loss of goodwill,
206
+ work stoppage, computer failure or malfunction, or any and all
207
+ other commercial damages or losses), even if such Contributor
208
+ has been advised of the possibility of such damages.
209
+
210
+ 9. Accepting Warranty or Additional Liability. While redistributing
211
+ the Work or Derivative Works thereof, You may choose to offer,
212
+ and charge a fee for, acceptance of support, warranty, indemnity,
213
+ or other liability obligations and/or rights consistent with this
214
+ License. However, in accepting such obligations, You may act only
215
+ on Your own behalf and on Your sole responsibility, not on behalf
216
+ of any other Contributor, and only if You agree to indemnify,
217
+ defend, and hold each Contributor harmless for any liability
218
+ incurred by, or claims asserted against, such Contributor by reason
219
+ of your accepting any such warranty or additional liability.
220
+
221
+ END OF TERMS AND CONDITIONS
222
+
223
+ APPENDIX: How to apply the Apache License to your work.
224
+
225
+ To apply the Apache License to your work, attach the following
226
+ boilerplate notice, with the fields enclosed by brackets "[]"
227
+ replaced with your own identifying information. (Don't include
228
+ the brackets!) The text should be enclosed in the appropriate
229
+ comment syntax for the file format. We also recommend that a
230
+ file or class name and description of purpose be included on the
231
+ same "printed page" as the copyright notice for easier
232
+ identification within third-party archives.
233
+
234
+ Copyright [yyyy] [name of copyright owner]
235
+
236
+ Licensed under the Apache License, Version 2.0 (the "License");
237
+ you may not use this file except in compliance with the License.
238
+ You may obtain a copy of the License at
239
+
240
+ http://www.apache.org/licenses/LICENSE-2.0
241
+
242
+ Unless required by applicable law or agreed to in writing, software
243
+ distributed under the License is distributed on an "AS IS" BASIS,
244
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
245
+ See the License for the specific language governing permissions and
246
+ limitations under the License.
inference/README.md ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # AX DocLayout v1.2 CUDA inference source
2
+
3
+ ONNX Runtime CUDA + TF32 inference for AXDocLayout-M v1.2, using unchanged FP32
4
+ weights and 1024 letterbox geometry. Requires supported Linux x86-64/NVIDIA
5
+ CUDA 12.8 setup and Python 3.12 or 3.13.
6
+
7
+ ```bash
8
+ python -m pip install ".[onnx]"
9
+ python examples/predict_page.py page.png --backend onnx
10
+ ```
11
+
12
+ The ONNX extra installs the pinned GPU runtime and CUDA/cuDNN dependencies.
13
+ The backend requires CUDA and verifies TF32/compute placement. A GPU failure
14
+ raises an error. Use the immutable source and weight revisions in the model
15
+ repository root README. For an explicit CPU comparison, install `.[torch]`
16
+ and select `--backend pytorch`. See LICENSE and NOTICE (AGPL-3.0-only).
inference/examples/predict_page.py ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Detect layout on one page raster using the pinned public Hub weights."""
5
+
6
+ from __future__ import annotations
7
+
8
+ import argparse
9
+ import json
10
+ import sys
11
+ import warnings
12
+ from pathlib import Path
13
+
14
+ import numpy as np
15
+ from PIL import Image
16
+
17
+ from ax_doclayout import AXDocLayout, AXDocLayoutError, InvalidPageImageError, PageImage
18
+ from ax_doclayout.contract import MAX_PAGE_PIXELS
19
+
20
+ MODEL_ID = "AutomatosX/AXDocLayout-M-v1.2"
21
+ REVISION = "7efd000fc7a1ef50ed7913835f2d8f10f30dc8e3"
22
+
23
+
24
+ def load_page(path: Path, *, page_index: int) -> PageImage:
25
+ """Decode a raster without recording its filename or pixels in the output."""
26
+
27
+ try:
28
+ with path.open("rb") as handle:
29
+ if handle.read(5) == b"%PDF-" or path.suffix.lower() == ".pdf":
30
+ raise InvalidPageImageError("provide a page raster image")
31
+ with warnings.catch_warnings():
32
+ warnings.simplefilter("error", Image.DecompressionBombWarning)
33
+ with Image.open(path) as image:
34
+ if image.width * image.height > MAX_PAGE_PIXELS:
35
+ raise InvalidPageImageError("page raster exceeds pixel limit")
36
+ rgb = np.asarray(image.convert("RGB"), dtype=np.uint8)
37
+ except (OSError, ValueError, Image.DecompressionBombError, Image.DecompressionBombWarning):
38
+ raise InvalidPageImageError("page raster could not be decoded") from None
39
+ return PageImage(rgb=rgb, page_index=page_index)
40
+
41
+
42
+ def main(argv: list[str] | None = None) -> int:
43
+ parser = argparse.ArgumentParser(description=__doc__)
44
+ parser.add_argument("image", type=Path)
45
+ parser.add_argument("--backend", choices=("onnx", "pytorch"), default="onnx")
46
+ parser.add_argument("--page-index", type=int, default=0)
47
+ args = parser.parse_args(argv)
48
+ try:
49
+ page = load_page(args.image, page_index=args.page_index)
50
+ model = AXDocLayout.from_pretrained(
51
+ MODEL_ID, revision=REVISION, backend=args.backend,
52
+ )
53
+ print(json.dumps(model.predict(page).to_json(), indent=2, allow_nan=False))
54
+ except AXDocLayoutError as error:
55
+ print(f"ax-doclayout: {error}", file=sys.stderr)
56
+ return 2
57
+ return 0
58
+
59
+
60
+ if __name__ == "__main__":
61
+ raise SystemExit(main())
inference/pyproject.toml ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [build-system]
2
+ requires = ["setuptools>=77", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "ax-doclayout"
7
+ version = "1.2.0.post2"
8
+ description = "AX DocLayout page-raster layout inference."
9
+ readme = "README.md"
10
+ requires-python = ">=3.12,<3.14"
11
+ license = "AGPL-3.0-only"
12
+ license-files = ["LICENSE", "NOTICE"]
13
+ authors = [{ name = "AutomatosX" }]
14
+ dependencies = ["huggingface-hub>=0.36,<2", "numpy>=1.26,<3", "pillow>=12.2,<13"]
15
+
16
+ [project.optional-dependencies]
17
+ onnx = [
18
+ "onnxruntime-gpu[cuda,cudnn]==1.23.2",
19
+ "nvidia-cuda-runtime-cu12==12.8.90",
20
+ "nvidia-cuda-nvrtc-cu12==12.8.93",
21
+ "nvidia-cufft-cu12==11.3.3.83",
22
+ "nvidia-curand-cu12==10.3.9.90",
23
+ "nvidia-cublas-cu12==12.8.4.1",
24
+ "nvidia-cudnn-cu12==9.10.2.21",
25
+ ]
26
+ torch = ["torch>=2.12,<2.13", "safetensors>=0.4,<1"]
27
+
28
+ [project.scripts]
29
+ ax-doclayout = "ax_doclayout.cli:main"
30
+
31
+ [tool.setuptools.packages.find]
32
+ where = ["src"]
33
+
34
+ [tool.setuptools.package-data]
35
+ ax_doclayout = ["py.typed"]
inference/src/ax_doclayout/__init__.py ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """AX DocLayout: page-raster document structure detection.
5
+
6
+ The AXDocLayout-M v1.2 model contract is a layout detector: an RGB page
7
+ raster in, versioned JSON boxes out. The public API verifies the pinned
8
+ immutable Hub bundle before inference. ``LocalOnnxBundle`` is an
9
+ unqualified local runner pinned to the staged ONNX digest; it is not Hub load.
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ from ax_doclayout.artifact import SOURCE_CHECKPOINT_SHA256, STAGED_ONNX_SHA256
15
+ from ax_doclayout.contract import (
16
+ COORDINATE_SPACE,
17
+ DETECTION_SCHEMA_VERSION,
18
+ PREPROCESS_ID,
19
+ TAXONOMY_ID,
20
+ BoundingBox,
21
+ Element,
22
+ InvalidPageStructureError,
23
+ ModelIdentity,
24
+ PageImage,
25
+ PageStructure,
26
+ )
27
+ from ax_doclayout.errors import (
28
+ ArtifactUnavailableError,
29
+ AXDocLayoutError,
30
+ BackendUnavailableError,
31
+ InferenceError,
32
+ InvalidPageImageError,
33
+ InvalidRevisionError,
34
+ )
35
+ from ax_doclayout.local_onnx import QUALIFICATION, LocalOnnxBundle, LocalOnnxResult
36
+ from ax_doclayout.model import AXDocLayout
37
+ from ax_doclayout.postprocess import (
38
+ DEFAULT_SCORE_THRESHOLD,
39
+ elements_from_letterbox_detections,
40
+ page_structure_from_elements,
41
+ )
42
+ from ax_doclayout.preprocess import (
43
+ Letterbox,
44
+ inverse_letterbox_box,
45
+ letterbox_page,
46
+ letterbox_rgb,
47
+ )
48
+
49
+ __version__ = "1.2.0.post2"
50
+
51
+ __all__ = [
52
+ "COORDINATE_SPACE",
53
+ "DEFAULT_SCORE_THRESHOLD",
54
+ "DETECTION_SCHEMA_VERSION",
55
+ "PREPROCESS_ID",
56
+ "QUALIFICATION",
57
+ "SOURCE_CHECKPOINT_SHA256",
58
+ "STAGED_ONNX_SHA256",
59
+ "TAXONOMY_ID",
60
+ "AXDocLayout",
61
+ "AXDocLayoutError",
62
+ "ArtifactUnavailableError",
63
+ "BackendUnavailableError",
64
+ "BoundingBox",
65
+ "Element",
66
+ "InferenceError",
67
+ "InvalidPageImageError",
68
+ "InvalidPageStructureError",
69
+ "InvalidRevisionError",
70
+ "Letterbox",
71
+ "LocalOnnxBundle",
72
+ "LocalOnnxResult",
73
+ "ModelIdentity",
74
+ "PageImage",
75
+ "PageStructure",
76
+ "__version__",
77
+ "elements_from_letterbox_detections",
78
+ "inverse_letterbox_box",
79
+ "letterbox_page",
80
+ "letterbox_rgb",
81
+ "page_structure_from_elements",
82
+ ]
inference/src/ax_doclayout/_cuda_onnx.py ADDED
@@ -0,0 +1,115 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Verified ONNX Runtime CUDA execution with explicit TF32 policy."""
5
+
6
+ from __future__ import annotations
7
+
8
+ import ctypes
9
+ import importlib.metadata
10
+ import json
11
+ import tempfile
12
+ from collections import Counter
13
+ from pathlib import Path
14
+ from typing import Any
15
+
16
+ import numpy as np
17
+
18
+ from ax_doclayout.errors import BackendUnavailableError
19
+
20
+ _COMPUTE_OPS = {"Conv", "FusedConv", "MatMul", "FusedMatMul", "Gemm", "GridSample", "Softmax"}
21
+ _PROVIDER = "CUDAExecutionProvider"
22
+ _OPTIONS = {"use_tf32": "1", "cudnn_conv_algo_search": "HEURISTIC", "device_id": "0"}
23
+
24
+
25
+ def _placement(events: object) -> dict[str, dict[str, int]]:
26
+ """Require observed CUDA compute and retain all remaining CPU placements."""
27
+ if not isinstance(events, list):
28
+ raise BackendUnavailableError("cuda_placement_unverified")
29
+ providers: Counter[str] = Counter()
30
+ cpu_ops: Counter[str] = Counter()
31
+ compute: Counter[str] = Counter()
32
+ for event in events:
33
+ if not isinstance(event, dict) or not isinstance(event.get("args"), dict):
34
+ continue
35
+ args = event["args"]
36
+ provider, op = args.get("provider"), args.get("op_name")
37
+ if not isinstance(provider, str) or not isinstance(op, str):
38
+ continue
39
+ providers[provider] += 1
40
+ if provider != _PROVIDER:
41
+ cpu_ops[op] += 1
42
+ if op in _COMPUTE_OPS:
43
+ if provider != _PROVIDER:
44
+ raise BackendUnavailableError("cuda_compute_placement_failed")
45
+ compute[op] += 1
46
+ if not any("Conv" in op for op in compute):
47
+ raise BackendUnavailableError("cuda_compute_unverified")
48
+ return {"providers": dict(providers), "cpu_ops": dict(cpu_ops), "cuda_compute": dict(compute)}
49
+
50
+
51
+ def _session(path: Path, *, ort: Any, profile_prefix: str | None = None) -> Any:
52
+ options = ort.SessionOptions()
53
+ options.intra_op_num_threads = 4
54
+ options.inter_op_num_threads = 1
55
+ options.log_severity_level = 3
56
+ if profile_prefix is not None:
57
+ options.enable_profiling = True
58
+ options.profile_file_prefix = profile_prefix
59
+ session = ort.InferenceSession(
60
+ str(path), sess_options=options,
61
+ providers=[(_PROVIDER, dict(_OPTIONS)), "CPUExecutionProvider"],
62
+ )
63
+ if not session.get_providers() or session.get_providers()[0] != _PROVIDER:
64
+ raise BackendUnavailableError("cuda_provider_not_primary")
65
+ session.disable_fallback()
66
+ actual = session.get_provider_options().get(_PROVIDER, {})
67
+ if actual.get("use_tf32") != "1" or actual.get("device_id") != "0":
68
+ raise BackendUnavailableError("cuda_tf32_policy_unverified")
69
+ return session
70
+
71
+
72
+ def _require_tf32_device() -> None:
73
+ """Check CUDA-visible device zero supports TF32 without importing PyTorch."""
74
+ distribution = importlib.metadata.distribution("nvidia-cuda-runtime-cu12")
75
+ library = distribution.locate_file("nvidia/cuda_runtime/lib/libcudart.so.12")
76
+ runtime = ctypes.CDLL(str(library))
77
+ query = runtime.cudaDeviceGetAttribute
78
+ query.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_int]
79
+ query.restype = ctypes.c_int
80
+ major = ctypes.c_int()
81
+ # CUDA Runtime enum cudaDevAttrComputeCapabilityMajor is 75.
82
+ status = query(ctypes.byref(major), 75, 0)
83
+ if status != 0 or major.value < 8:
84
+ raise BackendUnavailableError("tf32_capable_cuda_device_required")
85
+
86
+
87
+ def create_cuda_session(path: Path, *, ort: Any) -> Any:
88
+ """Load only the CUDA TF32 path; reject unavailable or unverified execution.
89
+
90
+ The startup trace uses a synthetic raster and is deleted before returning.
91
+ CPU shape/control operations remain explicit; recognized compute must use
92
+ CUDA. This check does not claim that every graph node runs on the GPU.
93
+ """
94
+ try:
95
+ # Load the CUDA/cuDNN libraries installed by the pinned GPU extra.
96
+ if hasattr(ort, "preload_dlls"):
97
+ ort.preload_dlls(directory="")
98
+ if _PROVIDER not in ort.get_available_providers():
99
+ raise BackendUnavailableError("cuda_provider_unavailable")
100
+ _require_tf32_device()
101
+ with tempfile.TemporaryDirectory(prefix="ax-cuda-probe-") as directory:
102
+ probe = _session(path, ort=ort, profile_prefix=str(Path(directory) / "profile"))
103
+ outputs = probe.run(["labels", "boxes", "scores"], {
104
+ "images": np.zeros((1, 3, 1024, 1024), dtype=np.float32),
105
+ "orig_target_sizes": np.array([[1024, 1024]], dtype=np.int64),
106
+ })
107
+ if len(outputs) != 3 or not all(np.isfinite(value).all() for value in outputs):
108
+ raise BackendUnavailableError("cuda_probe_outputs_invalid")
109
+ _placement(json.loads(Path(probe.end_profiling()).read_text()))
110
+ del probe
111
+ return _session(path, ort=ort)
112
+ except BackendUnavailableError:
113
+ raise
114
+ except Exception: # noqa: BLE001 - redact third-party paths and library errors
115
+ raise BackendUnavailableError("cuda_backend_load_failed") from None
inference/src/ax_doclayout/_hub_pin.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ """Verified immutable identity of the AXDocLayout-M v1.2 bundle."""
4
+ from __future__ import annotations
5
+
6
+ from ax_doclayout.artifact import ArtifactFile, ArtifactReceipt
7
+
8
+ PINNED_HUB_RECEIPT = ArtifactReceipt(
9
+ model_id='AutomatosX/AXDocLayout-M-v1.2',
10
+ model_revision='7efd000fc7a1ef50ed7913835f2d8f10f30dc8e3',
11
+ artifact_sha256='sha256:bcbfd65dc33bc08c05c5c511fe52c22dd439baf04c3f09c64f15dd3a78e9782d',
12
+ artifact_files=(
13
+ ArtifactFile(name='LICENSE', sha256='sha256:0d96a4ff68ad6d4b6f1f30f713b18d5184912ba8dd389f86aa7710db079abcb0'),
14
+ ArtifactFile(name='MODEL_CARD.md', sha256='sha256:f9680372775ec966971cdb373fb73413e9f5306573ec72e480358c4f7717f5ec'),
15
+ ArtifactFile(name='NOTICE', sha256='sha256:9faf3d395b2e23fde5419f4bf6ac22c093da2f92a48ba7533e88259de0173317'),
16
+ ArtifactFile(name='artifact-manifest.json', sha256='sha256:823c64a31559fa80f113e6c3d7be7163c43049a0a14c40028f6080a4bd5175ac'),
17
+ ArtifactFile(name='config.json', sha256='sha256:c60877aed03fc6a53d16303e1c6f0d2c86fddaf49b6f77c84daa243b61412a85'),
18
+ ArtifactFile(name='model.onnx', sha256='sha256:2d5928719c391e424a1132f8c5eadb23e4d164d3576f00ff14010ebe55aaea81'),
19
+ ArtifactFile(name='model.safetensors', sha256='sha256:3b97ba1c0b1b0816f45e79a76bb2fd495e026d2394471b76ce1a442baa2989c4'),
20
+ ArtifactFile(name='preprocessor_config.json', sha256='sha256:96c6b8cb4afc4b1a6dec51ce44a24d68ea05b42cdf2b78ce1df12a0ef3ec9f5b'),
21
+ ArtifactFile(name='taxonomy.json', sha256='sha256:896360d61d534175ea9d13164c96021bfb2b8aca6936a587a3dc1197b8deec35'),
22
+ ),
23
+ )
inference/src/ax_doclayout/_local_onnx.py ADDED
@@ -0,0 +1,175 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Private, digest-verified ONNX inference over a local nine-file bundle.
5
+
6
+ This path is not ``from_pretrained``. It never claims a Hub revision, never
7
+ marks a bundle publishable, and never emits a DocLayout-YOLO result shape.
8
+ Callers inject ``session_run`` in tests so the default suite does not need
9
+ onnxruntime, CUDA, or the 127 MiB staging graph.
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ from collections.abc import Callable, Mapping
15
+ from dataclasses import dataclass
16
+ from pathlib import Path
17
+ from typing import Any
18
+
19
+ import numpy as np
20
+
21
+ from ax_doclayout.artifact import (
22
+ MODEL_ID,
23
+ STAGED_ONNX_SHA256,
24
+ verify_artifact_manifest,
25
+ )
26
+ from ax_doclayout.contract import Element, PageImage
27
+ from ax_doclayout.errors import (
28
+ ArtifactError,
29
+ ArtifactUnavailableError,
30
+ BackendUnavailableError,
31
+ )
32
+ from ax_doclayout.postprocess import (
33
+ DEFAULT_SCORE_THRESHOLD,
34
+ elements_from_letterbox_detections,
35
+ )
36
+ from ax_doclayout.preprocess import letterbox_nchw, letterbox_page, orig_target_sizes
37
+
38
+ QUALIFICATION = "local-bundle-unqualified"
39
+ ONNX_INPUT_IMAGES = "images"
40
+ ONNX_INPUT_SIZES = "orig_target_sizes"
41
+
42
+
43
+ SessionRun = Callable[[Mapping[str, np.ndarray]], object]
44
+
45
+
46
+ @dataclass(frozen=True, slots=True)
47
+ class LocalOnnxResult:
48
+ """One unqualified local-bundle detection. Not a Hub ``PageStructure``."""
49
+
50
+ page_index: int
51
+ width_px: int
52
+ height_px: int
53
+ elements: tuple[Element, ...]
54
+ qualification: str
55
+ bundle_root: str
56
+ threshold: float
57
+
58
+ def as_published_page_structure(self) -> None:
59
+ """Refuse to impersonate a Hub-qualified ``PageStructure``."""
60
+
61
+ raise ArtifactUnavailableError(
62
+ "local ONNX bundle is not a published AX DocLayout artifact"
63
+ )
64
+
65
+
66
+ class LocalOnnxBundle:
67
+ """Run the staged ONNX file after the manifest binds every payload digest."""
68
+
69
+ def __init__(
70
+ self,
71
+ directory: Path,
72
+ *,
73
+ session_run: SessionRun | None = None,
74
+ expected_onnx_sha256: str | None = STAGED_ONNX_SHA256,
75
+ ) -> None:
76
+ manifest = verify_artifact_manifest(directory)
77
+ if manifest.model_id != MODEL_ID:
78
+ raise ArtifactError("artifact_model_id_invalid")
79
+ onnx_digest = ""
80
+ for item in manifest.artifact_files:
81
+ if item.name == "model.onnx":
82
+ onnx_digest = item.sha256
83
+ break
84
+ if not onnx_digest:
85
+ raise ArtifactError("artifact_file_unavailable")
86
+ if (
87
+ expected_onnx_sha256 is not None
88
+ and onnx_digest != expected_onnx_sha256
89
+ ):
90
+ raise ArtifactError("artifact_onnx_digest_mismatch")
91
+ self._directory = directory
92
+ self._onnx_path = directory / "model.onnx"
93
+ if not self._onnx_path.is_file():
94
+ raise ArtifactError("artifact_file_unavailable")
95
+ self._root = _file_digest(directory / "artifact-manifest.json")
96
+ self._session_run = session_run
97
+
98
+ def predict(
99
+ self,
100
+ page: PageImage,
101
+ *,
102
+ threshold: float = DEFAULT_SCORE_THRESHOLD,
103
+ ) -> LocalOnnxResult:
104
+ """Detect on one RGB page. Output stays unqualified."""
105
+
106
+ letterbox = letterbox_page(page)
107
+ feeds = {
108
+ ONNX_INPUT_IMAGES: letterbox_nchw(letterbox),
109
+ ONNX_INPUT_SIZES: orig_target_sizes(letterbox),
110
+ }
111
+ labels, boxes, scores = _split_outputs(self._run(feeds))
112
+ elements = elements_from_letterbox_detections(
113
+ labels=labels,
114
+ boxes=boxes,
115
+ scores=scores,
116
+ page=page,
117
+ letterbox=letterbox,
118
+ threshold=threshold,
119
+ )
120
+ return LocalOnnxResult(
121
+ page_index=page.page_index,
122
+ width_px=page.width_px,
123
+ height_px=page.height_px,
124
+ elements=elements,
125
+ qualification=QUALIFICATION,
126
+ bundle_root=self._root,
127
+ threshold=float(threshold),
128
+ )
129
+
130
+ def _run(self, feeds: Mapping[str, np.ndarray]) -> object:
131
+ if self._session_run is not None:
132
+ return self._session_run(feeds)
133
+ try:
134
+ import onnxruntime as ort
135
+ except ImportError as error:
136
+ raise BackendUnavailableError("onnxruntime is not installed") from error
137
+ try:
138
+ session = ort.InferenceSession(
139
+ str(self._onnx_path),
140
+ providers=["CPUExecutionProvider"],
141
+ )
142
+ return session.run(None, dict(feeds))
143
+ except Exception as error: # onnxruntime exposes no common base error
144
+ raise ArtifactError("onnx_inference_failed") from error
145
+
146
+
147
+ def _split_outputs(outputs: object) -> tuple[Any, Any, Any]:
148
+ if not isinstance(outputs, (tuple, list)) or len(outputs) < 3:
149
+ raise ArtifactError("onnx_outputs_invalid")
150
+ return outputs[0], outputs[1], outputs[2]
151
+
152
+
153
+ def _file_digest(path: Path) -> str:
154
+ """Hash one regular file without following symlinks."""
155
+
156
+ import hashlib
157
+ import os
158
+
159
+ flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0)
160
+ try:
161
+ descriptor = os.open(path, flags)
162
+ except OSError as error:
163
+ raise ArtifactError("artifact_manifest_unavailable") from error
164
+ digest = hashlib.sha256()
165
+ try:
166
+ while True:
167
+ chunk = os.read(descriptor, 1024 * 1024)
168
+ if not chunk:
169
+ break
170
+ digest.update(chunk)
171
+ except OSError as error:
172
+ raise ArtifactError("artifact_manifest_unavailable") from error
173
+ finally:
174
+ os.close(descriptor)
175
+ return digest.hexdigest()
inference/src/ax_doclayout/_local_pytorch.py ADDED
@@ -0,0 +1,122 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Digest-verified PyTorch inference over a local nine-file bundle.
5
+
6
+ This path loads the staged safetensors into the in-repo RT-DETRv2-M
7
+ architecture. It is not ``from_pretrained`` and it never marks a bundle
8
+ publishable.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ from pathlib import Path
14
+
15
+ import torch
16
+
17
+ from ax_doclayout._torch_rtdetrv2 import build_rtdetrv2_r50vd_m
18
+ from ax_doclayout.artifact import (
19
+ MODEL_ID,
20
+ STAGED_SAFETENSORS_SHA256,
21
+ _reject_json_constant,
22
+ _unique_object,
23
+ verify_artifact_manifest,
24
+ )
25
+ from ax_doclayout.contract import Element, PageImage
26
+ from ax_doclayout.errors import ArtifactError, BackendUnavailableError
27
+ from ax_doclayout.postprocess import (
28
+ DEFAULT_SCORE_THRESHOLD,
29
+ elements_from_letterbox_detections,
30
+ )
31
+ from ax_doclayout.preprocess import letterbox_nchw, letterbox_page, orig_target_sizes
32
+
33
+ ARCHITECTURE_ID = "rtdetrv2_r50vd_m"
34
+
35
+
36
+ class LocalPytorchBundle:
37
+ """Run the staged safetensors after the manifest binds every payload digest."""
38
+
39
+ def __init__(
40
+ self,
41
+ directory: Path,
42
+ *,
43
+ expected_safetensors_sha256: str | None = STAGED_SAFETENSORS_SHA256,
44
+ ) -> None:
45
+ try:
46
+ from safetensors import SafetensorError
47
+ from safetensors.torch import load_file
48
+ except ImportError as error:
49
+ raise BackendUnavailableError("safetensors is not installed") from error
50
+ manifest = verify_artifact_manifest(directory)
51
+ if manifest.model_id != MODEL_ID:
52
+ raise ArtifactError("artifact_model_id_invalid")
53
+ weights_digest = ""
54
+ for item in manifest.artifact_files:
55
+ if item.name == "model.safetensors":
56
+ weights_digest = item.sha256
57
+ break
58
+ if not weights_digest:
59
+ raise ArtifactError("artifact_file_unavailable")
60
+ if (
61
+ expected_safetensors_sha256 is not None
62
+ and weights_digest != expected_safetensors_sha256
63
+ ):
64
+ raise ArtifactError("artifact_safetensors_digest_mismatch")
65
+ weights_path = directory / "model.safetensors"
66
+ if not weights_path.is_file():
67
+ raise ArtifactError("artifact_file_unavailable")
68
+ try:
69
+ state = load_file(str(weights_path))
70
+ except (SafetensorError, OSError, ValueError) as error:
71
+ raise ArtifactError("artifact_safetensors_unreadable") from error
72
+ model = build_rtdetrv2_r50vd_m()
73
+ try:
74
+ model.load_state_dict(state, strict=True)
75
+ except RuntimeError as error:
76
+ raise ArtifactError("artifact_state_dict_mismatch") from error
77
+ model.eval()
78
+ self._model = model
79
+ self._directory = directory
80
+
81
+ def predict_elements(
82
+ self,
83
+ page: PageImage,
84
+ *,
85
+ threshold: float = DEFAULT_SCORE_THRESHOLD,
86
+ ) -> tuple[Element, ...]:
87
+ """Detect on one RGB page and return original-raster elements."""
88
+
89
+ letterbox = letterbox_page(page)
90
+ images = torch.from_numpy(letterbox_nchw(letterbox))
91
+ sizes = torch.from_numpy(orig_target_sizes(letterbox))
92
+ with torch.no_grad():
93
+ labels, boxes, scores = self._model(images, sizes)
94
+ return elements_from_letterbox_detections(
95
+ labels=labels.cpu().numpy(),
96
+ boxes=boxes.cpu().numpy(),
97
+ scores=scores.cpu().numpy(),
98
+ page=page,
99
+ letterbox=letterbox,
100
+ threshold=threshold,
101
+ )
102
+
103
+
104
+ def read_publishable(path: Path) -> bool:
105
+ """Read the staging receipt flag. Missing or malformed receipts fail closed."""
106
+
107
+ import json
108
+
109
+ try:
110
+ payload = json.loads(
111
+ path.read_text(encoding="utf-8"),
112
+ object_pairs_hook=_unique_object,
113
+ parse_constant=_reject_json_constant,
114
+ )
115
+ except (OSError, ValueError) as error:
116
+ raise ArtifactError("staging_receipt_unreadable") from error
117
+ if not isinstance(payload, dict) or "publishable" not in payload:
118
+ raise ArtifactError("staging_publishable_missing")
119
+ flag = payload["publishable"]
120
+ if not isinstance(flag, bool):
121
+ raise ArtifactError("staging_publishable_invalid")
122
+ return flag
inference/src/ax_doclayout/_torch_rtdetrv2.py ADDED
@@ -0,0 +1,490 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Original RT-DETRv2-M modules for the staged AXDocLayout-M v1.2 safetensors.
5
+
6
+ The layer names follow the staged checkpoint so ``load_state_dict`` is strict.
7
+ This module does not open a Hub revision and does not serialize JSON.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import torch
13
+ from torch import nn
14
+ from torch.nn import functional as F
15
+
16
+ HIDDEN = 256
17
+ NUM_CLASSES = 12
18
+ NUM_QUERIES = 300
19
+ # The staged ONNX graph executes the first three decoder layers and the
20
+ # third pair of score and box heads. Later checkpoint layers stay loaded.
21
+ INFERENCE_LAYERS = 3
22
+ NUM_LEVELS = 3
23
+ NUM_HEADS = 8
24
+ NUM_POINTS = 4
25
+ OFFSET_SCALE = 0.5
26
+ FEAT_CHANNELS = (512, 1024, 2048)
27
+ STAGE_CHANNELS = (256, 512, 1024, 2048)
28
+ STAGE_BLOCKS = (3, 4, 6, 3)
29
+ STAGE_STRIDES = (1, 2, 2, 2)
30
+
31
+
32
+ def inverse_sigmoid(value: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
33
+ """Map probabilities back to logits with a finite clamp."""
34
+
35
+ clipped = value.clamp(0.0, 1.0)
36
+ numerator = clipped.clamp(min=eps)
37
+ denominator = (1.0 - clipped).clamp(min=eps)
38
+ return torch.log(numerator / denominator)
39
+
40
+
41
+ def sincos_position_embedding(
42
+ height: int, width: int, *, device: torch.device, dtype: torch.dtype
43
+ ) -> torch.Tensor:
44
+ """2D sine-cosine embedding, shape ``[1, height * width, hidden]``."""
45
+
46
+ grid_y, grid_x = torch.meshgrid(
47
+ torch.arange(height, device=device, dtype=dtype),
48
+ torch.arange(width, device=device, dtype=dtype),
49
+ indexing="ij",
50
+ )
51
+ omega_count = HIDDEN // 4
52
+ omega = torch.arange(omega_count, device=device, dtype=dtype) / omega_count
53
+ omega = 1.0 / (10000.0 ** omega)
54
+ out_y = grid_y.reshape(-1, 1) * omega.reshape(1, -1)
55
+ out_x = grid_x.reshape(-1, 1) * omega.reshape(1, -1)
56
+ return torch.cat((out_y.sin(), out_y.cos(), out_x.sin(), out_x.cos()), dim=1)[
57
+ None
58
+ ]
59
+
60
+
61
+ class ConvNorm(nn.Module):
62
+ """Convolution, batch norm, and an optional pointwise activation."""
63
+
64
+ def __init__(
65
+ self,
66
+ in_channels: int,
67
+ out_channels: int,
68
+ kernel: int,
69
+ stride: int,
70
+ activation: str | None,
71
+ ) -> None:
72
+ super().__init__()
73
+ self.conv = nn.Conv2d(
74
+ in_channels,
75
+ out_channels,
76
+ kernel,
77
+ stride=stride,
78
+ padding=kernel // 2,
79
+ bias=False,
80
+ )
81
+ self.norm = nn.BatchNorm2d(out_channels)
82
+ self.activation = activation
83
+
84
+ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
85
+ tensor = self.norm(self.conv(tensor))
86
+ if self.activation == "relu":
87
+ return F.relu(tensor)
88
+ if self.activation == "silu":
89
+ return F.silu(tensor)
90
+ return tensor
91
+
92
+
93
+ class StrideShortcut(nn.Module):
94
+ """Average-pool then a stride-1 projection. Used by stage downsamples."""
95
+
96
+ def __init__(self, in_channels: int, out_channels: int) -> None:
97
+ super().__init__()
98
+ self.pool = nn.AvgPool2d(kernel_size=2, stride=2, ceil_mode=True)
99
+ self.conv = ConvNorm(in_channels, out_channels, 1, 1, None)
100
+
101
+ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
102
+ return self.conv(self.pool(tensor))
103
+
104
+
105
+ class ResidualBlock(nn.Module):
106
+ """Bottleneck with a ReLU after the residual add."""
107
+
108
+ def __init__(self, in_channels: int, out_channels: int, stride: int) -> None:
109
+ super().__init__()
110
+ mid = out_channels // 4
111
+ self.branch2a = ConvNorm(in_channels, mid, 1, 1, "relu")
112
+ self.branch2b = ConvNorm(mid, mid, 3, stride, "relu")
113
+ self.branch2c = ConvNorm(mid, out_channels, 1, 1, None)
114
+ if stride == 2:
115
+ self.short = StrideShortcut(in_channels, out_channels)
116
+ elif in_channels != out_channels:
117
+ self.short = ConvNorm(in_channels, out_channels, 1, 1, None)
118
+ else:
119
+ self.short = None
120
+
121
+ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
122
+ residual = tensor if self.short is None else self.short(tensor)
123
+ transformed = self.branch2c(self.branch2b(self.branch2a(tensor)))
124
+ return F.relu(transformed + residual)
125
+
126
+
127
+ class Backbone(nn.Module):
128
+ """PResNet-50-vd stem and four stages. Returns strides 8, 16, and 32."""
129
+
130
+ def __init__(self) -> None:
131
+ super().__init__()
132
+ self.conv1 = nn.Module()
133
+ self.conv1.conv1_1 = ConvNorm(3, 32, 3, 2, "relu")
134
+ self.conv1.conv1_2 = ConvNorm(32, 32, 3, 1, "relu")
135
+ self.conv1.conv1_3 = ConvNorm(32, 64, 3, 1, "relu")
136
+ self.pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
137
+ stages = []
138
+ in_channels = 64
139
+ for out_channels, blocks, stride in zip(
140
+ STAGE_CHANNELS, STAGE_BLOCKS, STAGE_STRIDES, strict=True
141
+ ):
142
+ layers = []
143
+ for index in range(blocks):
144
+ block_stride = stride if index == 0 else 1
145
+ layers.append(ResidualBlock(in_channels, out_channels, block_stride))
146
+ in_channels = out_channels
147
+ stage = nn.Module()
148
+ stage.blocks = nn.ModuleList(layers)
149
+ stages.append(stage)
150
+ self.res_layers = nn.ModuleList(stages)
151
+
152
+ def forward(self, tensor: torch.Tensor) -> list[torch.Tensor]:
153
+ tensor = self.conv1.conv1_3(self.conv1.conv1_2(self.conv1.conv1_1(tensor)))
154
+ tensor = self.pool(tensor)
155
+ outputs: list[torch.Tensor] = []
156
+ for stage_index, stage in enumerate(self.res_layers):
157
+ for block in stage.blocks:
158
+ tensor = block(tensor)
159
+ if stage_index > 0:
160
+ outputs.append(tensor)
161
+ return outputs
162
+
163
+
164
+ class RepVggBlock(nn.Module):
165
+ """3x3 plus 1x1 branches, summed, then SiLU."""
166
+
167
+ def __init__(self, channels: int) -> None:
168
+ super().__init__()
169
+ self.conv1 = ConvNorm(channels, channels, 3, 1, None)
170
+ self.conv2 = ConvNorm(channels, channels, 1, 1, None)
171
+
172
+ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
173
+ return F.silu(self.conv1(tensor) + self.conv2(tensor))
174
+
175
+
176
+ class CspRepBlock(nn.Module):
177
+ """Two 1x1 paths, repeated RepVGG blocks, then a 1x1 projection."""
178
+
179
+ def __init__(self, in_channels: int, out_channels: int, blocks: int = 3) -> None:
180
+ super().__init__()
181
+ hidden = out_channels // 2
182
+ self.conv1 = ConvNorm(in_channels, hidden, 1, 1, "silu")
183
+ self.conv2 = ConvNorm(in_channels, hidden, 1, 1, "silu")
184
+ self.bottlenecks = nn.Sequential(
185
+ *[RepVggBlock(hidden) for _ in range(blocks)]
186
+ )
187
+ self.conv3 = ConvNorm(hidden, out_channels, 1, 1, "silu")
188
+
189
+ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
190
+ return self.conv3(self.bottlenecks(self.conv1(tensor)) + self.conv2(tensor))
191
+
192
+
193
+ class EncoderLayer(nn.Module):
194
+ """One post-norm self-attention block with a GELU feed-forward."""
195
+
196
+ def __init__(self) -> None:
197
+ super().__init__()
198
+ self.self_attn = nn.MultiheadAttention(HIDDEN, NUM_HEADS, batch_first=True)
199
+ self.linear1 = nn.Linear(HIDDEN, HIDDEN * 4)
200
+ self.linear2 = nn.Linear(HIDDEN * 4, HIDDEN)
201
+ self.norm1 = nn.LayerNorm(HIDDEN)
202
+ self.norm2 = nn.LayerNorm(HIDDEN)
203
+
204
+ def forward(self, tensor: torch.Tensor, position: torch.Tensor) -> torch.Tensor:
205
+ query = tensor + position
206
+ attended, _ = self.self_attn(query, query, tensor, need_weights=False)
207
+ tensor = self.norm1(tensor + attended)
208
+ return self.norm2(tensor + self.linear2(F.gelu(self.linear1(tensor))))
209
+
210
+
211
+ class HybridEncoder(nn.Module):
212
+ """Project three backbone levels, encode the deepest, then FPN and PAN."""
213
+
214
+ def __init__(self) -> None:
215
+ super().__init__()
216
+ self.input_proj = nn.ModuleList(
217
+ [ConvNorm(channels, HIDDEN, 1, 1, None) for channels in FEAT_CHANNELS]
218
+ )
219
+ layer = EncoderLayer()
220
+ tower = nn.Module()
221
+ tower.layers = nn.ModuleList([layer])
222
+ self.encoder = nn.ModuleList([tower])
223
+ self.lateral_convs = nn.ModuleList(
224
+ [ConvNorm(HIDDEN, HIDDEN, 1, 1, "silu") for _ in range(2)]
225
+ )
226
+ self.fpn_blocks = nn.ModuleList(
227
+ [CspRepBlock(HIDDEN * 2, HIDDEN) for _ in range(2)]
228
+ )
229
+ self.downsample_convs = nn.ModuleList(
230
+ [ConvNorm(HIDDEN, HIDDEN, 3, 2, "silu") for _ in range(2)]
231
+ )
232
+ self.pan_blocks = nn.ModuleList(
233
+ [CspRepBlock(HIDDEN * 2, HIDDEN) for _ in range(2)]
234
+ )
235
+
236
+ def forward(self, features: list[torch.Tensor]) -> list[torch.Tensor]:
237
+ projected = [
238
+ layer(feature) for layer, feature in zip(self.input_proj, features, strict=True)
239
+ ]
240
+ deepest = projected[-1]
241
+ batch, _channels, height, width = deepest.shape
242
+ tokens = deepest.flatten(2).transpose(1, 2)
243
+ position = sincos_position_embedding(
244
+ height, width, device=tokens.device, dtype=tokens.dtype
245
+ )
246
+ tokens = self.encoder[0].layers[0](tokens, position)
247
+ projected[-1] = tokens.transpose(1, 2).reshape(batch, HIDDEN, height, width)
248
+
249
+ inner = [projected[-1]]
250
+ for index in range(NUM_LEVELS - 1, 0, -1):
251
+ block_index = NUM_LEVELS - 1 - index
252
+ high = self.lateral_convs[block_index](inner[0])
253
+ inner[0] = high
254
+ upsampled = F.interpolate(high, scale_factor=2.0, mode="nearest")
255
+ fused = self.fpn_blocks[block_index](
256
+ torch.cat((upsampled, projected[index - 1]), dim=1)
257
+ )
258
+ inner.insert(0, fused)
259
+
260
+ outputs = [inner[0]]
261
+ for index in range(NUM_LEVELS - 1):
262
+ downsampled = self.downsample_convs[index](outputs[-1])
263
+ outputs.append(
264
+ self.pan_blocks[index](torch.cat((downsampled, inner[index + 1]), dim=1))
265
+ )
266
+ return outputs
267
+
268
+
269
+ class MultiScaleAttention(nn.Module):
270
+ """Deformable attention over three feature levels, four points per head."""
271
+
272
+ def __init__(self) -> None:
273
+ super().__init__()
274
+ points = NUM_HEADS * NUM_LEVELS * NUM_POINTS
275
+ self.sampling_offsets = nn.Linear(HIDDEN, points * 2)
276
+ self.attention_weights = nn.Linear(HIDDEN, points)
277
+ self.value_proj = nn.Linear(HIDDEN, HIDDEN)
278
+ self.output_proj = nn.Linear(HIDDEN, HIDDEN)
279
+ self.num_points_scale = nn.Parameter(torch.empty(NUM_LEVELS * NUM_POINTS))
280
+
281
+ def forward(
282
+ self,
283
+ query: torch.Tensor,
284
+ reference: torch.Tensor,
285
+ value: torch.Tensor,
286
+ spatial_shapes: list[tuple[int, int]],
287
+ ) -> torch.Tensor:
288
+ batch, queries, _channels = query.shape
289
+ head_dim = HIDDEN // NUM_HEADS
290
+ projected = self.value_proj(value).view(batch, -1, NUM_HEADS, head_dim)
291
+ offsets = self.sampling_offsets(query).view(
292
+ batch, queries, NUM_HEADS, NUM_LEVELS, NUM_POINTS, 2
293
+ )
294
+ weights = self.attention_weights(query).view(
295
+ batch, queries, NUM_HEADS, NUM_LEVELS * NUM_POINTS
296
+ )
297
+ weights = torch.softmax(weights, dim=-1).view(
298
+ batch, queries, NUM_HEADS, NUM_LEVELS, NUM_POINTS
299
+ )
300
+ scale = self.num_points_scale.view(1, 1, 1, NUM_LEVELS, NUM_POINTS, 1)
301
+ locations = (
302
+ reference[:, :, None, :, None, :2]
303
+ + offsets * scale * reference[:, :, None, :, None, 2:] * OFFSET_SCALE
304
+ )
305
+ grids = locations.mul(2.0).sub(1.0)
306
+ split_sizes = [height * width for height, width in spatial_shapes]
307
+ values = projected.split(split_sizes, dim=1)
308
+ sampled = []
309
+ for level, ((height, width), level_value) in enumerate(
310
+ zip(spatial_shapes, values, strict=True)
311
+ ):
312
+ level_grid = (
313
+ grids[:, :, :, level]
314
+ .permute(0, 2, 1, 3, 4)
315
+ .reshape(batch * NUM_HEADS, queries, NUM_POINTS, 2)
316
+ )
317
+ image = (
318
+ level_value.reshape(batch, height, width, NUM_HEADS, head_dim)
319
+ .permute(0, 3, 4, 1, 2)
320
+ .reshape(batch * NUM_HEADS, head_dim, height, width)
321
+ )
322
+ sampled.append(
323
+ F.grid_sample(
324
+ image,
325
+ level_grid,
326
+ mode="bilinear",
327
+ padding_mode="zeros",
328
+ align_corners=False,
329
+ )
330
+ )
331
+ stacked = torch.stack(sampled, dim=-2)
332
+ # sampled: [B*heads, dim, queries, points] -> [B, Q, heads, levels, points, dim]
333
+ stacked = stacked.view(batch, NUM_HEADS, head_dim, queries, NUM_LEVELS, NUM_POINTS)
334
+ stacked = stacked.permute(0, 3, 1, 4, 5, 2)
335
+ weighted = (stacked * weights.unsqueeze(-1)).sum(dim=(3, 4))
336
+ merged = weighted.reshape(batch, queries, HIDDEN)
337
+ return self.output_proj(merged)
338
+
339
+
340
+ class DecoderLayer(nn.Module):
341
+ """Self-attention, multi-scale attention, and a ReLU feed-forward."""
342
+
343
+ def __init__(self) -> None:
344
+ super().__init__()
345
+ self.self_attn = nn.MultiheadAttention(HIDDEN, NUM_HEADS, batch_first=True)
346
+ self.cross_attn = MultiScaleAttention()
347
+ self.linear1 = nn.Linear(HIDDEN, HIDDEN * 4)
348
+ self.linear2 = nn.Linear(HIDDEN * 4, HIDDEN)
349
+ self.norm1 = nn.LayerNorm(HIDDEN)
350
+ self.norm2 = nn.LayerNorm(HIDDEN)
351
+ self.norm3 = nn.LayerNorm(HIDDEN)
352
+
353
+ def forward(
354
+ self,
355
+ target: torch.Tensor,
356
+ query_pos: torch.Tensor,
357
+ reference: torch.Tensor,
358
+ memory: torch.Tensor,
359
+ spatial_shapes: list[tuple[int, int]],
360
+ ) -> torch.Tensor:
361
+ query = target + query_pos
362
+ attended, _ = self.self_attn(query, query, target, need_weights=False)
363
+ target = self.norm1(target + attended)
364
+ cross = self.cross_attn(target + query_pos, reference, memory, spatial_shapes)
365
+ target = self.norm2(target + cross)
366
+ return self.norm3(target + self.linear2(F.relu(self.linear1(target))))
367
+
368
+
369
+ class Mlp(nn.Module):
370
+ """Linear stack with ReLU between layers and none on the last layer."""
371
+
372
+ def __init__(self, in_features: int, hidden: int, out_features: int, layers: int) -> None:
373
+ super().__init__()
374
+ widths = [in_features] + [hidden] * (layers - 1) + [out_features]
375
+ self.layers = nn.ModuleList(
376
+ nn.Linear(left, right) for left, right in zip(widths, widths[1:])
377
+ )
378
+
379
+ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
380
+ for index, layer in enumerate(self.layers):
381
+ tensor = layer(tensor)
382
+ if index + 1 < len(self.layers):
383
+ tensor = F.relu(tensor)
384
+ return tensor
385
+
386
+
387
+ class EncOutput(nn.Module):
388
+ """Project encoder memory and normalize it before the proposal heads."""
389
+
390
+ def __init__(self) -> None:
391
+ super().__init__()
392
+ self.proj = nn.Linear(HIDDEN, HIDDEN)
393
+ self.norm = nn.LayerNorm(HIDDEN)
394
+
395
+ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
396
+ return self.norm(self.proj(tensor))
397
+
398
+
399
+ class RTDETRV2R50VDM(nn.Module):
400
+ """Staged AXDocLayout-M v1.2 detector. Returns letterbox-space labels, boxes, and scores."""
401
+
402
+ def __init__(self) -> None:
403
+ super().__init__()
404
+ self.backbone = Backbone()
405
+ self.encoder = HybridEncoder()
406
+ self.decoder = nn.Module()
407
+ self.decoder.input_proj = nn.ModuleList(
408
+ [ConvNorm(HIDDEN, HIDDEN, 1, 1, None) for _ in range(NUM_LEVELS)]
409
+ )
410
+ self.decoder.enc_output = EncOutput()
411
+ self.decoder.enc_score_head = nn.Linear(HIDDEN, NUM_CLASSES)
412
+ self.decoder.enc_bbox_head = Mlp(HIDDEN, HIDDEN, 4, 3)
413
+ self.decoder.query_pos_head = Mlp(4, 512, HIDDEN, 2)
414
+ self.decoder.denoising_class_embed = nn.Embedding(NUM_CLASSES + 1, HIDDEN)
415
+ tower = nn.Module()
416
+ tower.layers = nn.ModuleList(DecoderLayer() for _ in range(6))
417
+ self.decoder.decoder = tower
418
+ self.decoder.dec_score_head = nn.ModuleList(
419
+ nn.Linear(HIDDEN, NUM_CLASSES) for _ in range(6)
420
+ )
421
+ self.decoder.dec_bbox_head = nn.ModuleList(
422
+ Mlp(HIDDEN, HIDDEN, 4, 3) for _ in range(6)
423
+ )
424
+ self.decoder.register_buffer("anchors", torch.empty(1, 21504, 4), persistent=True)
425
+ self.decoder.register_buffer(
426
+ "valid_mask", torch.empty(1, 21504, 1, dtype=torch.bool), persistent=True
427
+ )
428
+
429
+ def forward(
430
+ self, images: torch.Tensor, orig_target_sizes: torch.Tensor
431
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
432
+ features = self.encoder(self.backbone(images))
433
+ memories = []
434
+ shapes: list[tuple[int, int]] = []
435
+ for projection, feature in zip(self.decoder.input_proj, features, strict=True):
436
+ projected = projection(feature)
437
+ _batch, _channels, height, width = projected.shape
438
+ shapes.append((height, width))
439
+ memories.append(projected.flatten(2).transpose(1, 2))
440
+ memory = torch.cat(memories, dim=1)
441
+ valid = self.decoder.valid_mask.to(dtype=memory.dtype)
442
+ encoded = self.decoder.enc_output(memory * valid)
443
+ proposal_scores = self.decoder.enc_score_head(encoded).max(dim=-1).values
444
+ proposal_scores = proposal_scores.masked_fill(
445
+ ~self.decoder.valid_mask[..., 0], float("-inf")
446
+ )
447
+ _values, indices = torch.topk(proposal_scores, NUM_QUERIES, dim=1)
448
+ boxes = self.decoder.enc_bbox_head(encoded) + self.decoder.anchors
449
+ reference = torch.gather(
450
+ boxes, 1, indices.unsqueeze(-1).expand(-1, -1, 4)
451
+ ).detach().sigmoid()
452
+ target = torch.gather(
453
+ encoded, 1, indices.unsqueeze(-1).expand(-1, -1, HIDDEN)
454
+ ).detach()
455
+ level_reference = reference.unsqueeze(2).expand(-1, -1, NUM_LEVELS, -1)
456
+ logits = boxes_unactivated = None
457
+ for layer, score_head, box_head in zip(
458
+ list(self.decoder.decoder.layers)[:INFERENCE_LAYERS],
459
+ list(self.decoder.dec_score_head)[:INFERENCE_LAYERS],
460
+ list(self.decoder.dec_bbox_head)[:INFERENCE_LAYERS],
461
+ strict=True,
462
+ ):
463
+ query_pos = self.decoder.query_pos_head(reference)
464
+ target = layer(target, query_pos, level_reference, memory, shapes)
465
+ logits = score_head(target)
466
+ boxes_unactivated = box_head(target) + inverse_sigmoid(reference)
467
+ reference = boxes_unactivated.detach().sigmoid()
468
+ level_reference = reference.unsqueeze(2).expand(-1, -1, NUM_LEVELS, -1)
469
+ if logits is None or boxes_unactivated is None:
470
+ raise RuntimeError("decoder_did_not_run")
471
+ cx, cy, width, height = boxes_unactivated.sigmoid().unbind(-1)
472
+ xyxy = torch.stack(
473
+ (cx - 0.5 * width, cy - 0.5 * height, cx + 0.5 * width, cy + 0.5 * height),
474
+ dim=-1,
475
+ )
476
+ sizes = orig_target_sizes.to(dtype=xyxy.dtype)
477
+ scale = torch.stack((sizes[:, 0], sizes[:, 1], sizes[:, 0], sizes[:, 1]), dim=-1)
478
+ xyxy = xyxy * scale[:, None, :]
479
+ flat_scores = torch.sigmoid(logits).flatten(1)
480
+ scores, flat_index = torch.topk(flat_scores, NUM_QUERIES, dim=1)
481
+ labels = flat_index.remainder(NUM_CLASSES)
482
+ query_index = flat_index.div(NUM_CLASSES, rounding_mode="floor")
483
+ selected = torch.gather(xyxy, 1, query_index.unsqueeze(-1).expand(-1, -1, 4))
484
+ return labels, selected, scores
485
+
486
+
487
+ def build_rtdetrv2_r50vd_m() -> RTDETRV2R50VDM:
488
+ """Construct the unloaded AXDocLayout-M v1.2 architecture."""
489
+
490
+ return RTDETRV2R50VDM()
inference/src/ax_doclayout/artifact.py ADDED
@@ -0,0 +1,563 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Immutable AXDocLayout-M v1.2 artifact receipts and bundle verification.
5
+
6
+ The Hub bundle contains nine files. ``artifact-manifest.json`` describes the
7
+ other eight payload files, while a separately trusted receipt binds all nine
8
+ files and their content root. Keeping the receipt outside the bundle avoids a
9
+ self-digest cycle.
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ import hashlib
15
+ import json
16
+ import os
17
+ import re
18
+ import stat
19
+ from collections.abc import Mapping, Sequence
20
+ from dataclasses import dataclass
21
+ from pathlib import Path
22
+ from typing import Any
23
+
24
+ from ax_doclayout.contract import PREPROCESS_ID, TAXONOMY_ID
25
+ from ax_doclayout.errors import ArtifactError
26
+
27
+ ARTIFACT_ROOT_SCHEMA = "ax-doclayout.artifact-root/1"
28
+ ARTIFACT_RECEIPT_SCHEMA = "ax-doclayout.artifact-receipt/1"
29
+ ARTIFACT_MANIFEST_SCHEMA = "ax-doclayout.artifact-manifest/1"
30
+ MODEL_ID = "AutomatosX/AXDocLayout-M-v1.2"
31
+ LEGACY_MODEL_ID = "AutomatosX/AXDocStruct-M-v0.1"
32
+ SOURCE_CHECKPOINT_SHA256 = (
33
+ "sha256:e4747f71d40e14c89c561c720c043ad2be5dadf68668f071983f3ef1111a70fa"
34
+ )
35
+ STAGED_ONNX_SHA256 = (
36
+ "sha256:2d5928719c391e424a1132f8c5eadb23e4d164d3576f00ff14010ebe55aaea81"
37
+ )
38
+ STAGED_SAFETENSORS_SHA256 = (
39
+ "sha256:3b97ba1c0b1b0816f45e79a76bb2fd495e026d2394471b76ce1a442baa2989c4"
40
+ )
41
+ SUPPORTED_ONNX_OPSET = 16
42
+ ARTIFACT_BUNDLE_FILES = (
43
+ "LICENSE",
44
+ "MODEL_CARD.md",
45
+ "NOTICE",
46
+ "artifact-manifest.json",
47
+ "config.json",
48
+ "model.onnx",
49
+ "model.safetensors",
50
+ "preprocessor_config.json",
51
+ "taxonomy.json",
52
+ )
53
+ ARTIFACT_MANIFEST_PAYLOAD_FILES = tuple(
54
+ name for name in ARTIFACT_BUNDLE_FILES if name != "artifact-manifest.json"
55
+ )
56
+
57
+ _MAX_RECEIPT_BYTES = 64 * 1024
58
+ _MAX_MANIFEST_BYTES = 64 * 1024
59
+ _REVISION_RE = re.compile(r"^[0-9a-f]{40}$")
60
+ _SHA256_RE = re.compile(r"^sha256:[0-9a-f]{64}$")
61
+
62
+
63
+ @dataclass(frozen=True, slots=True)
64
+ class ArtifactFile:
65
+ """One manifest-relative artifact file identity."""
66
+
67
+ name: str
68
+ sha256: str
69
+
70
+ def __post_init__(self) -> None:
71
+ if self.name not in ARTIFACT_BUNDLE_FILES:
72
+ raise ArtifactError("artifact_file_name_invalid")
73
+ if (
74
+ not isinstance(self.sha256, str)
75
+ or _SHA256_RE.fullmatch(self.sha256) is None
76
+ ):
77
+ raise ArtifactError("artifact_file_digest_invalid")
78
+
79
+ def to_json(self) -> dict[str, str]:
80
+ """Return the canonical JSON shape."""
81
+ return {"name": self.name, "sha256": self.sha256}
82
+
83
+
84
+ @dataclass(frozen=True, slots=True)
85
+ class ArtifactReceipt:
86
+ """Trusted, package-external identity for one immutable Hub bundle."""
87
+
88
+ model_id: str
89
+ model_revision: str
90
+ artifact_sha256: str
91
+ artifact_files: Sequence[ArtifactFile]
92
+ taxonomy: str = TAXONOMY_ID
93
+ preprocess: str = PREPROCESS_ID
94
+ schema_version: str = ARTIFACT_RECEIPT_SCHEMA
95
+
96
+ def __post_init__(self) -> None:
97
+ files = tuple(self.artifact_files)
98
+ object.__setattr__(self, "artifact_files", files)
99
+ if self.schema_version != ARTIFACT_RECEIPT_SCHEMA:
100
+ raise ArtifactError("artifact_receipt_schema_invalid")
101
+ if self.model_id != MODEL_ID:
102
+ raise ArtifactError("artifact_model_id_invalid")
103
+ if (
104
+ not isinstance(self.model_revision, str)
105
+ or _REVISION_RE.fullmatch(self.model_revision) is None
106
+ ):
107
+ raise ArtifactError("artifact_revision_invalid")
108
+ if self.taxonomy != TAXONOMY_ID:
109
+ raise ArtifactError("artifact_taxonomy_invalid")
110
+ if self.preprocess != PREPROCESS_ID:
111
+ raise ArtifactError("artifact_preprocess_invalid")
112
+ if (
113
+ not isinstance(self.artifact_sha256, str)
114
+ or _SHA256_RE.fullmatch(self.artifact_sha256) is None
115
+ ):
116
+ raise ArtifactError("artifact_root_digest_invalid")
117
+ _validate_complete_files(files)
118
+ expected_root = artifact_root_sha256(files)
119
+ if self.artifact_sha256 != expected_root:
120
+ raise ArtifactError("artifact_root_digest_mismatch")
121
+
122
+ def to_json(self) -> dict[str, object]:
123
+ """Return the strict receipt JSON shape."""
124
+ return {
125
+ "artifact_files": [item.to_json() for item in self.artifact_files],
126
+ "artifact_sha256": self.artifact_sha256,
127
+ "model_id": self.model_id,
128
+ "model_revision": self.model_revision,
129
+ "preprocess": self.preprocess,
130
+ "schema_version": self.schema_version,
131
+ "taxonomy": self.taxonomy,
132
+ }
133
+
134
+
135
+ @dataclass(frozen=True, slots=True)
136
+ class ArtifactManifest:
137
+ """Canonical in-bundle identity for the eight non-manifest payloads."""
138
+
139
+ artifact_files: Sequence[ArtifactFile]
140
+ onnx_opset: int = SUPPORTED_ONNX_OPSET
141
+ model_id: str = MODEL_ID
142
+ taxonomy: str = TAXONOMY_ID
143
+ preprocess: str = PREPROCESS_ID
144
+ schema_version: str = ARTIFACT_MANIFEST_SCHEMA
145
+
146
+ def __post_init__(self) -> None:
147
+ files = tuple(self.artifact_files)
148
+ object.__setattr__(self, "artifact_files", files)
149
+ if self.schema_version != ARTIFACT_MANIFEST_SCHEMA:
150
+ raise ArtifactError("artifact_manifest_schema_invalid")
151
+ if self.model_id != MODEL_ID:
152
+ raise ArtifactError("artifact_model_id_invalid")
153
+ if self.taxonomy != TAXONOMY_ID:
154
+ raise ArtifactError("artifact_taxonomy_invalid")
155
+ if self.preprocess != PREPROCESS_ID:
156
+ raise ArtifactError("artifact_preprocess_invalid")
157
+ if (
158
+ isinstance(self.onnx_opset, bool)
159
+ or not isinstance(self.onnx_opset, int)
160
+ or self.onnx_opset != SUPPORTED_ONNX_OPSET
161
+ ):
162
+ raise ArtifactError("artifact_onnx_opset_unsupported")
163
+ _validate_manifest_files(files)
164
+
165
+ def to_json(self) -> dict[str, object]:
166
+ """Return the strict manifest JSON shape."""
167
+ return {
168
+ "backends": {
169
+ "onnx": {
170
+ "file": "model.onnx",
171
+ "format": "onnx",
172
+ "opset": self.onnx_opset,
173
+ "precision": "fp32",
174
+ "standard_ops_only": True,
175
+ },
176
+ "pytorch": {
177
+ "file": "model.safetensors",
178
+ "format": "safetensors",
179
+ "precision": "fp32",
180
+ },
181
+ },
182
+ "files": [item.to_json() for item in self.artifact_files],
183
+ "model_id": self.model_id,
184
+ "preprocess": self.preprocess,
185
+ "schema_version": self.schema_version,
186
+ "taxonomy": self.taxonomy,
187
+ }
188
+
189
+
190
+ def canonical_json_bytes(payload: Mapping[str, object]) -> bytes:
191
+ """Serialize strict ASCII JSON without a trailing newline."""
192
+ try:
193
+ return json.dumps(
194
+ payload,
195
+ allow_nan=False,
196
+ ensure_ascii=True,
197
+ separators=(",", ":"),
198
+ sort_keys=True,
199
+ ).encode("ascii")
200
+ except (TypeError, ValueError, UnicodeError) as exc:
201
+ raise ArtifactError("artifact_json_invalid") from exc
202
+
203
+
204
+ def _validate_complete_files(files: Sequence[ArtifactFile]) -> None:
205
+ if any(not isinstance(item, ArtifactFile) for item in files):
206
+ raise ArtifactError("artifact_files_invalid")
207
+ names = tuple(item.name for item in files)
208
+ if len(names) != len(set(names)):
209
+ raise ArtifactError("artifact_file_duplicate")
210
+ if names != ARTIFACT_BUNDLE_FILES:
211
+ raise ArtifactError("artifact_file_set_invalid")
212
+
213
+
214
+ def _validate_manifest_files(files: Sequence[ArtifactFile]) -> None:
215
+ if any(not isinstance(item, ArtifactFile) for item in files):
216
+ raise ArtifactError("artifact_files_invalid")
217
+ names = tuple(item.name for item in files)
218
+ if len(names) != len(set(names)):
219
+ raise ArtifactError("artifact_file_duplicate")
220
+ if names != ARTIFACT_MANIFEST_PAYLOAD_FILES:
221
+ raise ArtifactError("artifact_manifest_file_set_invalid")
222
+
223
+
224
+ def artifact_root_sha256(files: Sequence[ArtifactFile]) -> str:
225
+ """Compute the content root over the sorted complete nine-file set."""
226
+ items = tuple(files)
227
+ _validate_complete_files(items)
228
+ payload = {
229
+ "files": [item.to_json() for item in items],
230
+ "schema_version": ARTIFACT_ROOT_SCHEMA,
231
+ }
232
+ return "sha256:" + hashlib.sha256(canonical_json_bytes(payload)).hexdigest()
233
+
234
+
235
+ def _open_regular_file(path: Path) -> int:
236
+ flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0)
237
+ try:
238
+ descriptor = os.open(path, flags)
239
+ except OSError as exc:
240
+ raise ArtifactError("artifact_file_unavailable") from exc
241
+ try:
242
+ opened = os.fstat(descriptor)
243
+ observed = path.lstat()
244
+ if (
245
+ not stat.S_ISREG(opened.st_mode)
246
+ or not stat.S_ISREG(observed.st_mode)
247
+ or (opened.st_dev, opened.st_ino) != (observed.st_dev, observed.st_ino)
248
+ ):
249
+ raise ArtifactError("artifact_file_unavailable")
250
+ except (OSError, ArtifactError):
251
+ os.close(descriptor)
252
+ raise
253
+ return descriptor
254
+
255
+
256
+ def _hash_file(path: Path) -> str:
257
+ digest = hashlib.sha256()
258
+ descriptor = _open_regular_file(path)
259
+ try:
260
+ with os.fdopen(descriptor, "rb") as handle:
261
+ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
262
+ digest.update(chunk)
263
+ except OSError as exc:
264
+ raise ArtifactError("artifact_file_unavailable") from exc
265
+ return "sha256:" + digest.hexdigest()
266
+
267
+
268
+ def _bundle_files(directory: Path) -> tuple[ArtifactFile, ...]:
269
+ try:
270
+ directory_mode = directory.lstat().st_mode
271
+ except OSError as exc:
272
+ raise ArtifactError("artifact_bundle_unavailable") from exc
273
+ if not stat.S_ISDIR(directory_mode):
274
+ raise ArtifactError("artifact_bundle_unavailable")
275
+ try:
276
+ names = tuple(sorted(path.name for path in directory.iterdir()))
277
+ except OSError as exc:
278
+ raise ArtifactError("artifact_bundle_unavailable") from exc
279
+ if names != ARTIFACT_BUNDLE_FILES:
280
+ raise ArtifactError("artifact_bundle_file_set_invalid")
281
+ files: list[ArtifactFile] = []
282
+ for name in ARTIFACT_BUNDLE_FILES:
283
+ path = directory / name
284
+ files.append(ArtifactFile(name=name, sha256=_hash_file(path)))
285
+ return tuple(files)
286
+
287
+
288
+ def _payload_files(directory: Path) -> tuple[ArtifactFile, ...]:
289
+ try:
290
+ directory_mode = directory.lstat().st_mode
291
+ except OSError as exc:
292
+ raise ArtifactError("artifact_bundle_unavailable") from exc
293
+ if not stat.S_ISDIR(directory_mode):
294
+ raise ArtifactError("artifact_bundle_unavailable")
295
+ try:
296
+ names = tuple(sorted(path.name for path in directory.iterdir()))
297
+ except OSError as exc:
298
+ raise ArtifactError("artifact_bundle_unavailable") from exc
299
+ allowed_sets = {
300
+ ARTIFACT_MANIFEST_PAYLOAD_FILES,
301
+ ARTIFACT_BUNDLE_FILES,
302
+ }
303
+ if names not in allowed_sets:
304
+ raise ArtifactError("artifact_bundle_file_set_invalid")
305
+ return tuple(
306
+ ArtifactFile(name=name, sha256=_hash_file(directory / name))
307
+ for name in ARTIFACT_MANIFEST_PAYLOAD_FILES
308
+ )
309
+
310
+
311
+ def build_artifact_manifest(
312
+ directory: Path,
313
+ *,
314
+ onnx_opset: int = SUPPORTED_ONNX_OPSET,
315
+ ) -> ArtifactManifest:
316
+ """Build the manifest from the exact eight payload files."""
317
+ return ArtifactManifest(
318
+ artifact_files=_payload_files(directory),
319
+ onnx_opset=onnx_opset,
320
+ )
321
+
322
+
323
+ def build_artifact_receipt(
324
+ directory: Path,
325
+ *,
326
+ model_revision: str,
327
+ ) -> ArtifactReceipt:
328
+ """Build a receipt after a complete bundle has an immutable Hub revision."""
329
+ verify_artifact_manifest(directory)
330
+ files = _bundle_files(directory)
331
+ return ArtifactReceipt(
332
+ model_id=MODEL_ID,
333
+ model_revision=model_revision,
334
+ artifact_sha256=artifact_root_sha256(files),
335
+ artifact_files=files,
336
+ )
337
+
338
+
339
+ def verify_artifact_bundle(directory: Path, receipt: ArtifactReceipt) -> None:
340
+ """Fail closed unless every staged byte matches a trusted receipt."""
341
+ if not isinstance(receipt, ArtifactReceipt):
342
+ raise ArtifactError("artifact_receipt_invalid")
343
+ verify_artifact_manifest(directory)
344
+ observed = _bundle_files(directory)
345
+ if observed != receipt.artifact_files:
346
+ raise ArtifactError("artifact_file_digest_mismatch")
347
+ if artifact_root_sha256(observed) != receipt.artifact_sha256:
348
+ raise ArtifactError("artifact_root_digest_mismatch")
349
+
350
+
351
+ def _reject_json_constant(_: str) -> None:
352
+ raise ValueError("non_finite_json")
353
+
354
+
355
+ def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
356
+ result: dict[str, Any] = {}
357
+ for key, value in pairs:
358
+ if key in result:
359
+ raise ValueError("duplicate_json_key")
360
+ result[key] = value
361
+ return result
362
+
363
+
364
+ def _exact_keys(
365
+ value: Mapping[str, Any],
366
+ expected: set[str],
367
+ code: str = "artifact_receipt_fields_invalid",
368
+ ) -> None:
369
+ if set(value) != expected:
370
+ raise ArtifactError(code)
371
+
372
+
373
+ def _parse_artifact_file(value: object) -> ArtifactFile:
374
+ if not isinstance(value, dict):
375
+ raise ArtifactError("artifact_file_invalid")
376
+ _exact_keys(value, {"name", "sha256"}, "artifact_file_invalid")
377
+ try:
378
+ return ArtifactFile(name=value["name"], sha256=value["sha256"])
379
+ except (KeyError, TypeError) as exc:
380
+ raise ArtifactError("artifact_file_invalid") from exc
381
+
382
+
383
+ def _validated_backends(value: object) -> int:
384
+ if not isinstance(value, dict):
385
+ raise ArtifactError("artifact_backends_invalid")
386
+ _exact_keys(value, {"onnx", "pytorch"}, "artifact_backends_invalid")
387
+ onnx = value["onnx"]
388
+ pytorch = value["pytorch"]
389
+ if not isinstance(onnx, dict) or not isinstance(pytorch, dict):
390
+ raise ArtifactError("artifact_backends_invalid")
391
+ _exact_keys(
392
+ onnx,
393
+ {"file", "format", "opset", "precision", "standard_ops_only"},
394
+ "artifact_backends_invalid",
395
+ )
396
+ _exact_keys(
397
+ pytorch,
398
+ {"file", "format", "precision"},
399
+ "artifact_backends_invalid",
400
+ )
401
+ opset = onnx["opset"]
402
+ if (
403
+ onnx["file"] != "model.onnx"
404
+ or onnx["format"] != "onnx"
405
+ or onnx["precision"] != "fp32"
406
+ or onnx["standard_ops_only"] is not True
407
+ or isinstance(opset, bool)
408
+ or not isinstance(opset, int)
409
+ or opset != SUPPORTED_ONNX_OPSET
410
+ or pytorch
411
+ != {
412
+ "file": "model.safetensors",
413
+ "format": "safetensors",
414
+ "precision": "fp32",
415
+ }
416
+ ):
417
+ raise ArtifactError("artifact_backends_invalid")
418
+ return opset
419
+
420
+
421
+ def parse_artifact_manifest(raw: bytes) -> ArtifactManifest:
422
+ """Parse canonical manifest bytes with exact backend and payload metadata."""
423
+ if len(raw) > _MAX_MANIFEST_BYTES:
424
+ raise ArtifactError("artifact_manifest_too_large")
425
+ try:
426
+ value = json.loads(
427
+ raw,
428
+ object_pairs_hook=_unique_object,
429
+ parse_constant=_reject_json_constant,
430
+ )
431
+ except (TypeError, ValueError, UnicodeError, json.JSONDecodeError) as exc:
432
+ raise ArtifactError("artifact_manifest_json_invalid") from exc
433
+ if not isinstance(value, dict):
434
+ raise ArtifactError("artifact_manifest_invalid")
435
+ _exact_keys(
436
+ value,
437
+ {
438
+ "backends",
439
+ "files",
440
+ "model_id",
441
+ "preprocess",
442
+ "schema_version",
443
+ "taxonomy",
444
+ },
445
+ "artifact_manifest_fields_invalid",
446
+ )
447
+ raw_files = value["files"]
448
+ if not isinstance(raw_files, list):
449
+ raise ArtifactError("artifact_files_invalid")
450
+ manifest = ArtifactManifest(
451
+ artifact_files=tuple(_parse_artifact_file(item) for item in raw_files),
452
+ onnx_opset=_validated_backends(value["backends"]),
453
+ model_id=value["model_id"],
454
+ taxonomy=value["taxonomy"],
455
+ preprocess=value["preprocess"],
456
+ schema_version=value["schema_version"],
457
+ )
458
+ if canonical_json_bytes(manifest.to_json()) != raw:
459
+ raise ArtifactError("artifact_manifest_not_canonical")
460
+ return manifest
461
+
462
+
463
+ def load_artifact_manifest(path: Path) -> ArtifactManifest:
464
+ """Load one bounded canonical in-bundle manifest."""
465
+ descriptor = _open_regular_file(path)
466
+ try:
467
+ with os.fdopen(descriptor, "rb") as handle:
468
+ raw = handle.read(_MAX_MANIFEST_BYTES + 1)
469
+ except OSError as exc:
470
+ raise ArtifactError("artifact_manifest_unavailable") from exc
471
+ if len(raw) > _MAX_MANIFEST_BYTES:
472
+ raise ArtifactError("artifact_manifest_too_large")
473
+ return parse_artifact_manifest(raw)
474
+
475
+
476
+ def verify_artifact_manifest(directory: Path) -> ArtifactManifest:
477
+ """Verify the manifest against the exact eight payload bytes."""
478
+ manifest = load_artifact_manifest(directory / "artifact-manifest.json")
479
+ observed = _payload_files(directory)
480
+ if tuple(manifest.artifact_files) != observed:
481
+ raise ArtifactError("artifact_manifest_digest_mismatch")
482
+ return manifest
483
+
484
+
485
+ def parse_artifact_receipt(raw: bytes) -> ArtifactReceipt:
486
+ """Parse canonical trusted-receipt bytes with duplicate-key rejection."""
487
+ if len(raw) > _MAX_RECEIPT_BYTES:
488
+ raise ArtifactError("artifact_receipt_too_large")
489
+ try:
490
+ value = json.loads(
491
+ raw,
492
+ object_pairs_hook=_unique_object,
493
+ parse_constant=_reject_json_constant,
494
+ )
495
+ except (TypeError, ValueError, UnicodeError, json.JSONDecodeError) as exc:
496
+ raise ArtifactError("artifact_receipt_json_invalid") from exc
497
+ if not isinstance(value, dict):
498
+ raise ArtifactError("artifact_receipt_invalid")
499
+ _exact_keys(
500
+ value,
501
+ {
502
+ "artifact_files",
503
+ "artifact_sha256",
504
+ "model_id",
505
+ "model_revision",
506
+ "preprocess",
507
+ "schema_version",
508
+ "taxonomy",
509
+ },
510
+ )
511
+ raw_files = value["artifact_files"]
512
+ if not isinstance(raw_files, list):
513
+ raise ArtifactError("artifact_files_invalid")
514
+ receipt = ArtifactReceipt(
515
+ model_id=value["model_id"],
516
+ model_revision=value["model_revision"],
517
+ artifact_sha256=value["artifact_sha256"],
518
+ artifact_files=tuple(_parse_artifact_file(item) for item in raw_files),
519
+ taxonomy=value["taxonomy"],
520
+ preprocess=value["preprocess"],
521
+ schema_version=value["schema_version"],
522
+ )
523
+ if canonical_json_bytes(receipt.to_json()) != raw:
524
+ raise ArtifactError("artifact_receipt_not_canonical")
525
+ return receipt
526
+
527
+
528
+ def load_trusted_receipt(path: Path) -> ArtifactReceipt:
529
+ """Load one operator-pinned local receipt without contacting the Hub."""
530
+ descriptor = _open_regular_file(path)
531
+ try:
532
+ with os.fdopen(descriptor, "rb") as handle:
533
+ raw = handle.read(_MAX_RECEIPT_BYTES + 1)
534
+ except OSError as exc:
535
+ raise ArtifactError("artifact_receipt_unavailable") from exc
536
+ if len(raw) > _MAX_RECEIPT_BYTES:
537
+ raise ArtifactError("artifact_receipt_too_large")
538
+ return parse_artifact_receipt(raw)
539
+
540
+
541
+ __all__ = [
542
+ "ARTIFACT_BUNDLE_FILES",
543
+ "ARTIFACT_MANIFEST_PAYLOAD_FILES",
544
+ "ARTIFACT_MANIFEST_SCHEMA",
545
+ "ARTIFACT_RECEIPT_SCHEMA",
546
+ "ARTIFACT_ROOT_SCHEMA",
547
+ "LEGACY_MODEL_ID",
548
+ "MODEL_ID",
549
+ "SUPPORTED_ONNX_OPSET",
550
+ "ArtifactFile",
551
+ "ArtifactManifest",
552
+ "ArtifactReceipt",
553
+ "artifact_root_sha256",
554
+ "build_artifact_manifest",
555
+ "build_artifact_receipt",
556
+ "canonical_json_bytes",
557
+ "load_artifact_manifest",
558
+ "load_trusted_receipt",
559
+ "parse_artifact_manifest",
560
+ "parse_artifact_receipt",
561
+ "verify_artifact_bundle",
562
+ "verify_artifact_manifest",
563
+ ]
inference/src/ax_doclayout/cli.py ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """``ax-doclayout`` command line: one page raster in, canonical JSON out."""
5
+
6
+ from __future__ import annotations
7
+
8
+ import argparse
9
+ import json
10
+ import sys
11
+ from pathlib import Path
12
+
13
+ from ax_doclayout.contract import PageImage
14
+ from ax_doclayout.errors import AXDocLayoutError, InvalidPageImageError
15
+ from ax_doclayout.model import BACKENDS, AXDocLayout
16
+
17
+ _PDF_MAGIC = b"%PDF-"
18
+
19
+
20
+ def _load_rgb(path: Path):
21
+ """Decode a page raster; PDFs are rejected by content and by file name."""
22
+
23
+ try:
24
+ with path.open("rb") as stream:
25
+ head = stream.read(len(_PDF_MAGIC))
26
+ except OSError as error:
27
+ raise InvalidPageImageError(
28
+ f"cannot read page raster: {error.strerror}"
29
+ ) from None
30
+ if head.startswith(_PDF_MAGIC) or path.suffix.lower() == ".pdf":
31
+ raise InvalidPageImageError("the CLI takes a decoded page raster, not a PDF")
32
+ import numpy as np
33
+ from PIL import Image, UnidentifiedImageError
34
+
35
+ try:
36
+ with Image.open(path) as image:
37
+ return np.asarray(image.convert("RGB"), dtype=np.uint8)
38
+ except (OSError, UnidentifiedImageError, Image.DecompressionBombError, ValueError) as error:
39
+ raise InvalidPageImageError(
40
+ f"cannot decode page raster: {type(error).__name__}"
41
+ ) from None
42
+
43
+
44
+ def main(argv: list[str] | None = None) -> int:
45
+ parser = argparse.ArgumentParser(prog="ax-doclayout")
46
+ commands = parser.add_subparsers(dest="command", required=True)
47
+ predict = commands.add_parser(
48
+ "predict", help="detect layout regions on one page raster"
49
+ )
50
+ predict.add_argument("image", type=Path)
51
+ predict.add_argument("--page-index", type=int, required=True)
52
+ predict.add_argument("--model", required=True)
53
+ predict.add_argument("--revision", required=True)
54
+ predict.add_argument("--backend", choices=BACKENDS, required=True)
55
+ predict.add_argument("--output", default="-")
56
+ arguments = parser.parse_args(argv)
57
+
58
+ try:
59
+ page = PageImage(
60
+ rgb=_load_rgb(arguments.image), page_index=arguments.page_index
61
+ )
62
+ model = AXDocLayout.from_pretrained(
63
+ arguments.model, revision=arguments.revision, backend=arguments.backend
64
+ )
65
+ payload = json.dumps(model.predict(page).to_json(), indent=2)
66
+ except AXDocLayoutError as error:
67
+ print(f"ax-doclayout: {error}", file=sys.stderr)
68
+ return 2
69
+ if arguments.output == "-":
70
+ print(payload)
71
+ else:
72
+ Path(arguments.output).write_text(payload + "\n", encoding="utf-8")
73
+ return 0
74
+
75
+
76
+ if __name__ == "__main__":
77
+ raise SystemExit(main())
inference/src/ax_doclayout/contract.py ADDED
@@ -0,0 +1,286 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Typed values of the raster detection contract (spec sections 4, 5, 7, 9).
5
+
6
+ Everything here is the consumer-facing shape; nothing here runs a model.
7
+ Every violation raises an ``AXDocLayoutError`` subclass.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import math
13
+ import re
14
+ from collections.abc import Mapping, Sequence
15
+ from dataclasses import dataclass, field
16
+ from numbers import Integral, Real
17
+ from typing import Any
18
+
19
+ import numpy as np
20
+
21
+ from ax_doclayout.errors import AXDocLayoutError, InvalidPageImageError
22
+
23
+ # stayed because the payload shape is unchanged. The previous prefix is
24
+ # not an alias.
25
+ DETECTION_SCHEMA_VERSION = "ax-doclayout.detection/1"
26
+ TAXONOMY_ID = "ax-doclayout-11/1"
27
+ PREPROCESS_ID = "rgb-letterbox-1024/1"
28
+ COORDINATE_SPACE = "raster-pixels-top-left"
29
+ DEFAULT_THRESHOLD_PROFILE = "default/1"
30
+
31
+ # Public class ids are preserved; id 0 is reserved and never returned.
32
+ CLASS_NAMES: Mapping[int, str] = {
33
+ 1: "caption",
34
+ 2: "footnote",
35
+ 3: "formula",
36
+ 4: "list_item",
37
+ 5: "page_footer",
38
+ 6: "page_header",
39
+ 7: "picture",
40
+ 8: "section_header",
41
+ 9: "table",
42
+ 10: "text",
43
+ 11: "title",
44
+ }
45
+
46
+ # Rasters past this many pixels are rejected before any preprocessing.
47
+ MAX_PAGE_PIXELS = 50_000_000
48
+
49
+ HUB_COMMIT_RE = re.compile(r"^[0-9a-f]{40}$")
50
+ SHA256_RE = re.compile(r"^[0-9a-f]{64}$")
51
+
52
+
53
+ class InvalidPageStructureError(AXDocLayoutError):
54
+ """A detection result violates the canonical contract."""
55
+
56
+
57
+ def _integer(value: object, name: str, *, minimum: int = 0) -> int:
58
+ """Return ``value`` as a plain int; bools and non-integral types are rejected."""
59
+
60
+ if isinstance(value, bool) or not isinstance(value, Integral):
61
+ raise InvalidPageStructureError(f"{name} must be an integer")
62
+ result = int(value)
63
+ if result < minimum:
64
+ raise InvalidPageStructureError(f"{name} must be >= {minimum}")
65
+ return result
66
+
67
+
68
+ def _finite(value: object, name: str) -> float:
69
+ if isinstance(value, bool) or not isinstance(value, Real):
70
+ raise InvalidPageStructureError(f"{name} must be a real number")
71
+ result = float(value)
72
+ if not math.isfinite(result):
73
+ raise InvalidPageStructureError(f"{name} must be finite")
74
+ return result
75
+
76
+
77
+ @dataclass(frozen=True, slots=True, eq=False)
78
+ class PageImage:
79
+ """One decoded page raster: HWC ``uint8`` RGB plus its zero-based page index.
80
+
81
+ The array is copied to C-contiguous, read-only form on construction so
82
+ later caller mutation can neither change a validated page nor be frozen
83
+ as a side effect of constructing one.
84
+ """
85
+
86
+ rgb: np.ndarray
87
+ page_index: int
88
+
89
+ def __post_init__(self) -> None:
90
+ if not isinstance(self.rgb, np.ndarray):
91
+ raise InvalidPageImageError("rgb must be a numpy array")
92
+ if self.rgb.ndim != 3 or self.rgb.shape[2] != 3:
93
+ raise InvalidPageImageError("rgb must be HWC with three channels")
94
+ if self.rgb.dtype != np.uint8:
95
+ raise InvalidPageImageError("rgb must be uint8")
96
+ height, width = self.rgb.shape[:2]
97
+ if height == 0 or width == 0:
98
+ raise InvalidPageImageError("rgb must not be empty")
99
+ if height * width > MAX_PAGE_PIXELS:
100
+ raise InvalidPageImageError("rgb exceeds the page pixel limit")
101
+ if isinstance(self.page_index, bool) or not isinstance(
102
+ self.page_index, Integral
103
+ ):
104
+ raise InvalidPageImageError("page_index must be an int")
105
+ if self.page_index < 0:
106
+ raise InvalidPageImageError(
107
+ "page_index must be zero-based and non-negative"
108
+ )
109
+ rgb = np.array(self.rgb, dtype=np.uint8, order="C", copy=True)
110
+ rgb.setflags(write=False)
111
+ object.__setattr__(self, "rgb", rgb)
112
+ object.__setattr__(self, "page_index", int(self.page_index))
113
+
114
+ @property
115
+ def width_px(self) -> int:
116
+ return int(self.rgb.shape[1])
117
+
118
+ @property
119
+ def height_px(self) -> int:
120
+ return int(self.rgb.shape[0])
121
+
122
+
123
+ @dataclass(frozen=True, slots=True)
124
+ class BoundingBox:
125
+ """Half-open pixel-edge box ``[x_min, x_max) x [y_min, y_max)`` on the original raster."""
126
+
127
+ x_min: float
128
+ y_min: float
129
+ x_max: float
130
+ y_max: float
131
+
132
+ def validate(self, width_px: int, height_px: int) -> None:
133
+ x_min = _finite(self.x_min, "x_min")
134
+ y_min = _finite(self.y_min, "y_min")
135
+ x_max = _finite(self.x_max, "x_max")
136
+ y_max = _finite(self.y_max, "y_max")
137
+ if not (0 <= x_min < x_max <= width_px):
138
+ raise InvalidPageStructureError(
139
+ "bounding box x range must satisfy 0 <= x_min < x_max <= width_px"
140
+ )
141
+ if not (0 <= y_min < y_max <= height_px):
142
+ raise InvalidPageStructureError(
143
+ "bounding box y range must satisfy 0 <= y_min < y_max <= height_px"
144
+ )
145
+
146
+ def to_json(self) -> dict[str, float]:
147
+ return {
148
+ "x_min": float(self.x_min),
149
+ "y_min": float(self.y_min),
150
+ "x_max": float(self.x_max),
151
+ "y_max": float(self.y_max),
152
+ }
153
+
154
+
155
+ @dataclass(frozen=True, slots=True)
156
+ class Element:
157
+ """One detected region. ``id`` is opaque and stable only within its result."""
158
+
159
+ id: str
160
+ class_id: int
161
+ confidence: float
162
+ bbox: BoundingBox
163
+
164
+ @property
165
+ def class_name(self) -> str:
166
+ if isinstance(self.class_id, bool) or self.class_id not in CLASS_NAMES:
167
+ raise InvalidPageStructureError(f"unknown class id {self.class_id!r}")
168
+ return CLASS_NAMES[int(self.class_id)]
169
+
170
+ def to_json(self) -> dict[str, Any]:
171
+ return {
172
+ "id": self.id,
173
+ "class_id": int(self.class_id),
174
+ "class_name": self.class_name,
175
+ "confidence": float(self.confidence),
176
+ "bbox": self.bbox.to_json(),
177
+ }
178
+
179
+
180
+ @dataclass(frozen=True, slots=True)
181
+ class ModelIdentity:
182
+ """The immutable artifact a result came from (spec section 7)."""
183
+
184
+ id: str
185
+ revision: str
186
+ artifact_sha256: str
187
+
188
+ def __post_init__(self) -> None:
189
+ if not isinstance(self.id, str) or not self.id.strip():
190
+ raise InvalidPageStructureError("model id must be a non-empty string")
191
+ if (
192
+ not isinstance(self.revision, str)
193
+ or HUB_COMMIT_RE.fullmatch(self.revision) is None
194
+ ):
195
+ raise InvalidPageStructureError(
196
+ "model revision must be a 40-character immutable Hub commit"
197
+ )
198
+ if (
199
+ not isinstance(self.artifact_sha256, str)
200
+ or SHA256_RE.fullmatch(self.artifact_sha256) is None
201
+ ):
202
+ raise InvalidPageStructureError(
203
+ "artifact_sha256 must be a 64-character hex digest"
204
+ )
205
+
206
+ def to_json(self) -> dict[str, str]:
207
+ return {
208
+ "id": self.id,
209
+ "revision": self.revision,
210
+ "artifact_sha256": self.artifact_sha256,
211
+ "taxonomy": TAXONOMY_ID,
212
+ "preprocess": PREPROCESS_ID,
213
+ }
214
+
215
+
216
+ @dataclass(frozen=True, slots=True)
217
+ class PageStructure:
218
+ """The canonical detection result for one page (spec section 9)."""
219
+
220
+ model: ModelIdentity
221
+ page_index: int
222
+ width_px: int
223
+ height_px: int
224
+ elements: Sequence[Element] = field(default_factory=tuple)
225
+ threshold_profile: str = DEFAULT_THRESHOLD_PROFILE
226
+ threshold_override: float | None = None
227
+
228
+ def __post_init__(self) -> None:
229
+ if not isinstance(self.model, ModelIdentity):
230
+ raise InvalidPageStructureError("model must be a ModelIdentity")
231
+ object.__setattr__(self, "page_index", _integer(self.page_index, "page_index"))
232
+ object.__setattr__(
233
+ self, "width_px", _integer(self.width_px, "width_px", minimum=1)
234
+ )
235
+ object.__setattr__(
236
+ self, "height_px", _integer(self.height_px, "height_px", minimum=1)
237
+ )
238
+ if (
239
+ not isinstance(self.threshold_profile, str)
240
+ or not self.threshold_profile.strip()
241
+ ):
242
+ raise InvalidPageStructureError(
243
+ "threshold_profile must be a non-empty string"
244
+ )
245
+ if self.threshold_override is not None:
246
+ override = _finite(self.threshold_override, "threshold_override")
247
+ if not 0.0 <= override <= 1.0:
248
+ raise InvalidPageStructureError(
249
+ "threshold_override must be within [0, 1]"
250
+ )
251
+ object.__setattr__(self, "threshold_override", override)
252
+ elements = tuple(self.elements)
253
+ object.__setattr__(self, "elements", elements)
254
+ seen: set[str] = set()
255
+ for element in elements:
256
+ if not isinstance(element, Element):
257
+ raise InvalidPageStructureError("elements must be Element values")
258
+ if not isinstance(element.id, str) or not element.id:
259
+ raise InvalidPageStructureError("element ids must be non-empty strings")
260
+ if element.id in seen:
261
+ raise InvalidPageStructureError(f"duplicate element id {element.id!r}")
262
+ seen.add(element.id)
263
+ _ = element.class_name # raises for reserved or unknown ids
264
+ confidence = _finite(element.confidence, "confidence")
265
+ if not 0.0 <= confidence <= 1.0:
266
+ raise InvalidPageStructureError("confidence must be within [0, 1]")
267
+ if not isinstance(element.bbox, BoundingBox):
268
+ raise InvalidPageStructureError("element bbox must be a BoundingBox")
269
+ element.bbox.validate(self.width_px, self.height_px)
270
+
271
+ def to_json(self) -> dict[str, Any]:
272
+ inference: dict[str, Any] = {"threshold_profile": self.threshold_profile}
273
+ if self.threshold_override is not None:
274
+ inference["threshold_override"] = float(self.threshold_override)
275
+ return {
276
+ "schema_version": DETECTION_SCHEMA_VERSION,
277
+ "model": self.model.to_json(),
278
+ "page": {
279
+ "page_index": int(self.page_index),
280
+ "width_px": int(self.width_px),
281
+ "height_px": int(self.height_px),
282
+ "coordinate_space": COORDINATE_SPACE,
283
+ },
284
+ "inference": inference,
285
+ "elements": [element.to_json() for element in self.elements],
286
+ }
inference/src/ax_doclayout/errors.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Public error types. Messages stay content-free."""
5
+ from __future__ import annotations
6
+
7
+
8
+ class AXDocLayoutError(Exception):
9
+ """Base error for the public contract."""
10
+
11
+
12
+ class InvalidPageImageError(AXDocLayoutError):
13
+ """Page raster failed validation."""
14
+
15
+
16
+ class PdfRejectedError(AXDocLayoutError):
17
+ """CLI and loaders reject PDF inputs."""
18
+
19
+
20
+ class UnsupportedBackendError(AXDocLayoutError):
21
+ """Backend missing, unknown, or unavailable."""
22
+
23
+
24
+ class ArtifactError(AXDocLayoutError):
25
+ """Artifact identity, digest, or files failed closed."""
26
+
27
+
28
+ class ArtifactUnavailableError(ArtifactError):
29
+ """No published immutable artifact is available to load."""
30
+
31
+
32
+ class InvalidRevisionError(ArtifactError):
33
+ """Hub revision is not a 40-character commit."""
34
+
35
+
36
+ class BackendUnavailableError(UnsupportedBackendError):
37
+ """Requested inference backend is missing or unknown."""
38
+
39
+
40
+ class InferenceError(AXDocLayoutError):
41
+ """A loaded backend could not execute page-raster inference."""
42
+
43
+
44
+ class TaxonomyError(AXDocLayoutError):
45
+ """Unknown or reserved class identity."""
inference/src/ax_doclayout/hub_release.py ADDED
@@ -0,0 +1,205 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Trusted Hub identity and local verification of the FP32 model bundle."""
5
+
6
+ from __future__ import annotations
7
+
8
+ import json
9
+ import os
10
+ from pathlib import Path
11
+
12
+ from ax_doclayout import _hub_pin
13
+ from ax_doclayout.artifact import (
14
+ MODEL_ID,
15
+ STAGED_ONNX_SHA256,
16
+ STAGED_SAFETENSORS_SHA256,
17
+ ArtifactReceipt,
18
+ _bundle_files,
19
+ _open_regular_file,
20
+ build_artifact_receipt,
21
+ verify_artifact_bundle,
22
+ verify_artifact_manifest,
23
+ )
24
+ from ax_doclayout.contract import HUB_COMMIT_RE
25
+ from ax_doclayout.errors import (
26
+ ArtifactError,
27
+ ArtifactUnavailableError,
28
+ BackendUnavailableError,
29
+ InvalidRevisionError,
30
+ )
31
+
32
+ PINNED_HUB_RECEIPT = _hub_pin.PINNED_HUB_RECEIPT
33
+
34
+ _MAX_SAFETENSORS_HEADER = 8 * 1024 * 1024
35
+ # FP32 weights plus the integer batch counters and the one bool mask that
36
+ # PyTorch stores beside them. BF16, FP16, and FP64 are rejected.
37
+ _ALLOWED_SAFETENSORS_DTYPES = frozenset({"F32", "I64", "BOOL"})
38
+
39
+
40
+ def safetensors_dtypes(path: Path) -> tuple[str, ...]:
41
+ """Return tensor dtypes from a safetensors header, without reading weights."""
42
+
43
+ descriptor = _open_regular_file(path)
44
+ try:
45
+ with os.fdopen(descriptor, "rb") as handle:
46
+ prefix = handle.read(8)
47
+ if len(prefix) != 8:
48
+ raise ArtifactError("staging_safetensors_header_invalid")
49
+ length = int.from_bytes(prefix, "little")
50
+ if length <= 0 or length > _MAX_SAFETENSORS_HEADER:
51
+ raise ArtifactError("staging_safetensors_header_invalid")
52
+ raw = handle.read(length)
53
+ except OSError as exc:
54
+ raise ArtifactError("staging_safetensors_unreadable") from exc
55
+ if len(raw) != length:
56
+ raise ArtifactError("staging_safetensors_header_invalid")
57
+ try:
58
+ header = json.loads(raw.decode("utf-8"))
59
+ except (UnicodeError, ValueError) as exc:
60
+ raise ArtifactError("staging_safetensors_header_invalid") from exc
61
+ if not isinstance(header, dict):
62
+ raise ArtifactError("staging_safetensors_header_invalid")
63
+ dtypes: list[str] = []
64
+ for key, value in header.items():
65
+ if not isinstance(key, str):
66
+ raise ArtifactError("staging_safetensors_header_invalid")
67
+ if key.startswith("__"):
68
+ continue
69
+ if not isinstance(value, dict) or not isinstance(value.get("dtype"), str):
70
+ raise ArtifactError("staging_safetensors_header_invalid")
71
+ dtypes.append(value["dtype"])
72
+ if not dtypes:
73
+ raise ArtifactError("staging_safetensors_empty")
74
+ return tuple(dtypes)
75
+
76
+
77
+ def _require_fp32_weights(path: Path) -> None:
78
+ dtypes = set(safetensors_dtypes(path))
79
+ if "F32" not in dtypes or not dtypes <= _ALLOWED_SAFETENSORS_DTYPES:
80
+ raise ArtifactError("staging_safetensors_not_fp32")
81
+
82
+
83
+ def verify_staged_v01_bundle(
84
+ directory: Path,
85
+ receipt_path: Path,
86
+ *,
87
+ require_staged_digests: bool = True,
88
+ ) -> None:
89
+ """Fail closed unless a nine-file bundle is FP32 and still unpublished."""
90
+
91
+ from ax_doclayout._local_pytorch import read_publishable
92
+
93
+ if read_publishable(receipt_path):
94
+ raise ArtifactError("staging_publishable_must_stay_false")
95
+ verify_artifact_manifest(directory)
96
+ observed = _bundle_files(directory)
97
+ by_name = {item.name: item.sha256 for item in observed}
98
+ if require_staged_digests and (
99
+ by_name["model.safetensors"] != STAGED_SAFETENSORS_SHA256
100
+ or by_name["model.onnx"] != STAGED_ONNX_SHA256
101
+ ):
102
+ raise ArtifactError("staging_weight_digest_mismatch")
103
+ _require_fp32_weights(directory / "model.safetensors")
104
+ _require_receipt_agreement(receipt_path, by_name)
105
+
106
+
107
+ def _require_receipt_agreement(receipt_path: Path, by_name: dict[str, str]) -> None:
108
+ try:
109
+ payload = json.loads(receipt_path.read_text(encoding="utf-8"))
110
+ except (OSError, UnicodeError, ValueError) as exc:
111
+ raise ArtifactError("staging_receipt_unreadable") from exc
112
+ if not isinstance(payload, dict):
113
+ raise ArtifactError("staging_receipt_unreadable")
114
+ model_id = payload.get("model_id", MODEL_ID)
115
+ if model_id != MODEL_ID:
116
+ raise ArtifactError("artifact_model_id_invalid")
117
+ files = payload.get("artifact_files")
118
+ if files is None:
119
+ return
120
+ if not isinstance(files, list):
121
+ raise ArtifactError("staging_receipt_unreadable")
122
+ for item in files:
123
+ if not isinstance(item, dict):
124
+ raise ArtifactError("staging_receipt_unreadable")
125
+ name = item.get("name")
126
+ digest = item.get("sha256")
127
+ if name in by_name and digest != by_name[name]:
128
+ raise ArtifactError("staging_receipt_digest_mismatch")
129
+
130
+
131
+ def qualify_published_bundle(
132
+ *,
133
+ model_id: str,
134
+ revision: str,
135
+ backend: str,
136
+ receipt: ArtifactReceipt | None,
137
+ bundle: Path | None,
138
+ ) -> Path | None:
139
+ """Verify one pinned receipt against local bytes.
140
+
141
+ ``None`` means no production receipt is pinned. A pinned receipt with a
142
+ missing or mismatched bundle raises. This function does not download and
143
+ does not run a backend.
144
+ """
145
+
146
+ if receipt is None:
147
+ return None
148
+ if backend not in ("pytorch", "onnx"):
149
+ raise BackendUnavailableError(
150
+ "backend must be given explicitly as one of pytorch, onnx"
151
+ )
152
+ if model_id != MODEL_ID or model_id != receipt.model_id:
153
+ raise ArtifactError("artifact_model_id_invalid")
154
+ if revision != receipt.model_revision:
155
+ raise ArtifactError("artifact_revision_invalid")
156
+ if bundle is None:
157
+ raise ArtifactUnavailableError(
158
+ "published bundle is not available for hash verification"
159
+ )
160
+ verify_artifact_bundle(bundle, receipt)
161
+ _require_fp32_weights(bundle / "model.safetensors")
162
+ return bundle
163
+
164
+
165
+ SIGN_OFF_FIELDS = ("adr_0031", "license_pack", "model_card")
166
+
167
+
168
+ def require_publication_signoff(path: Path) -> None:
169
+ """Refuse publication until ADR-0031 and the license pack are approved."""
170
+
171
+ try:
172
+ payload = json.loads(path.read_text(encoding="utf-8"))
173
+ except (OSError, UnicodeError, ValueError) as exc:
174
+ raise ArtifactError("publication_signoff_missing") from exc
175
+ if not isinstance(payload, dict):
176
+ raise ArtifactError("publication_signoff_missing")
177
+ for field in SIGN_OFF_FIELDS:
178
+ if payload.get(field) != "approved":
179
+ raise ArtifactError("publication_signoff_incomplete")
180
+
181
+
182
+ def publish_redownloaded_bundle(
183
+ *,
184
+ bundle: Path,
185
+ staging_receipt: Path,
186
+ signoff: Path,
187
+ revision: str,
188
+ downloaded: Path,
189
+ require_staged_digests: bool = True,
190
+ ) -> ArtifactReceipt:
191
+ """Build a receipt only after a redownload matches the staged nine files."""
192
+
193
+ require_publication_signoff(signoff)
194
+ verify_staged_v01_bundle(
195
+ bundle,
196
+ staging_receipt,
197
+ require_staged_digests=require_staged_digests,
198
+ )
199
+ if not isinstance(revision, str) or HUB_COMMIT_RE.fullmatch(revision) is None:
200
+ raise InvalidRevisionError(
201
+ "revision must be a 40-character immutable Hub commit"
202
+ )
203
+ if _bundle_files(bundle) != _bundle_files(downloaded):
204
+ raise ArtifactError("hub_redownload_digest_mismatch")
205
+ return build_artifact_receipt(downloaded, model_revision=revision)
inference/src/ax_doclayout/local_detection.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Canonical JSON for unqualified local detections.
5
+
6
+ Both the local ONNX runner and the local PyTorch runner serialize through
7
+ this function. The payload is not a published ``PageStructure``.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from collections.abc import Sequence
13
+
14
+ from ax_doclayout.artifact import canonical_json_bytes
15
+ from ax_doclayout.contract import Element
16
+
17
+ LOCAL_DETECTION_SCHEMA = "ax-doclayout.local-detection/1"
18
+ LOCAL_QUALIFICATION = "local-bundle-unqualified"
19
+ # Shared unqualified serialization. PyTorch and ONNX agree well inside
20
+ # these places; the published parity gate is coarser (0.5 px, 0.001 score).
21
+ SCORE_DECIMALS = 4
22
+ BOX_DECIMALS = 3
23
+
24
+
25
+ def _places(value: float, places: int) -> float:
26
+ return round(float(value), places)
27
+
28
+
29
+ def canonical_local_detection_bytes(
30
+ *,
31
+ page_index: int,
32
+ width_px: int,
33
+ height_px: int,
34
+ elements: Sequence[Element],
35
+ threshold: float,
36
+ ) -> bytes:
37
+ """Serialize one unqualified page detection with sorted keys."""
38
+
39
+ payload = {
40
+ "elements": [
41
+ {
42
+ "bbox": {
43
+ "x_max": _places(element.bbox.x_max, BOX_DECIMALS),
44
+ "x_min": _places(element.bbox.x_min, BOX_DECIMALS),
45
+ "y_max": _places(element.bbox.y_max, BOX_DECIMALS),
46
+ "y_min": _places(element.bbox.y_min, BOX_DECIMALS),
47
+ },
48
+ "class_id": int(element.class_id),
49
+ "class_name": element.class_name,
50
+ "confidence": _places(element.confidence, SCORE_DECIMALS),
51
+ "id": element.id,
52
+ }
53
+ for element in elements
54
+ ],
55
+ "height_px": int(height_px),
56
+ "page_index": int(page_index),
57
+ "qualification": LOCAL_QUALIFICATION,
58
+ "schema_version": LOCAL_DETECTION_SCHEMA,
59
+ "threshold": float(threshold),
60
+ "width_px": int(width_px),
61
+ }
62
+ return canonical_json_bytes(payload)
inference/src/ax_doclayout/local_onnx.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Unqualified local nine-file ONNX runner.
5
+
6
+ This is a public operator surface for a digest-pinned staging bundle. It is
7
+ not ``from_pretrained``, does not claim a Hub revision, and does not emit a
8
+ DocLayout-YOLO result shape.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ from ax_doclayout._local_onnx import (
14
+ QUALIFICATION,
15
+ LocalOnnxBundle,
16
+ LocalOnnxResult,
17
+ )
18
+ from ax_doclayout.artifact import SOURCE_CHECKPOINT_SHA256, STAGED_ONNX_SHA256
19
+
20
+ __all__ = [
21
+ "QUALIFICATION",
22
+ "SOURCE_CHECKPOINT_SHA256",
23
+ "STAGED_ONNX_SHA256",
24
+ "LocalOnnxBundle",
25
+ "LocalOnnxResult",
26
+ ]
inference/src/ax_doclayout/model.py ADDED
@@ -0,0 +1,158 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Load a verified immutable Hub bundle and detect layout on RGB pages."""
5
+
6
+ from __future__ import annotations
7
+
8
+ import importlib
9
+ import shutil
10
+ from collections.abc import Callable
11
+ from pathlib import Path
12
+ from tempfile import TemporaryDirectory
13
+
14
+ from ax_doclayout.artifact import MODEL_ID, ArtifactReceipt
15
+ from ax_doclayout.contract import (
16
+ HUB_COMMIT_RE,
17
+ Element,
18
+ ModelIdentity,
19
+ PageImage,
20
+ PageStructure,
21
+ )
22
+ from ax_doclayout.errors import (
23
+ ArtifactError,
24
+ ArtifactUnavailableError,
25
+ AXDocLayoutError,
26
+ BackendUnavailableError,
27
+ InferenceError,
28
+ InvalidPageImageError,
29
+ InvalidRevisionError,
30
+ )
31
+ from ax_doclayout.hub_release import PINNED_HUB_RECEIPT, qualify_published_bundle
32
+ from ax_doclayout.postprocess import page_structure_from_elements
33
+
34
+ BACKENDS = ("pytorch", "onnx")
35
+ _PredictElements = Callable[[PageImage], tuple[Element, ...]]
36
+
37
+
38
+ def _download_bundle(directory: Path, *, receipt: ArtifactReceipt) -> None:
39
+ """Materialize regular files from the Hub cache before strict verification."""
40
+
41
+ try:
42
+ from huggingface_hub import hf_hub_download
43
+ except ImportError:
44
+ raise BackendUnavailableError("huggingface-hub is not installed") from None
45
+ try:
46
+ for item in receipt.artifact_files:
47
+ cached = hf_hub_download(
48
+ repo_id=receipt.model_id,
49
+ filename=item.name,
50
+ revision=receipt.model_revision,
51
+ token=False,
52
+ )
53
+ # Hub cache entries may be symlinks. Only the copied, hash-verified
54
+ # regular files are handed to artifact validation and the backend.
55
+ shutil.copyfile(cached, directory / item.name)
56
+ except Exception: # noqa: BLE001 - redact third-party download error details
57
+ raise ArtifactUnavailableError("hub_bundle_download_failed") from None
58
+
59
+
60
+ def _load_backend(directory: Path, *, backend: str) -> _PredictElements:
61
+ """Load all backend state before the temporary verified bundle is removed."""
62
+
63
+ if backend == "pytorch":
64
+ from ax_doclayout._local_pytorch import LocalPytorchBundle
65
+
66
+ return LocalPytorchBundle(directory).predict_elements
67
+
68
+ import onnxruntime as ort
69
+
70
+ from ax_doclayout._cuda_onnx import create_cuda_session
71
+ from ax_doclayout._local_onnx import LocalOnnxBundle
72
+
73
+ session = create_cuda_session(directory / "model.onnx", ort=ort)
74
+ runner = LocalOnnxBundle(
75
+ directory, session_run=lambda feeds: session.run(None, dict(feeds)),
76
+ )
77
+
78
+ def predict_elements(page: PageImage) -> tuple[Element, ...]:
79
+ return runner.predict(page).elements
80
+
81
+ return predict_elements
82
+
83
+
84
+ class AXDocLayout:
85
+ """A trusted Hub artifact with explicit CUDA ONNX or CPU PyTorch inference."""
86
+
87
+ def __init__(self, model_id: str, revision: str, backend: str) -> None:
88
+ self.model_id = model_id
89
+ self.revision = revision
90
+ self.backend = backend
91
+ self._predict_elements: _PredictElements | None = None
92
+ self._identity: ModelIdentity | None = None
93
+
94
+ @classmethod
95
+ def from_pretrained(
96
+ cls,
97
+ model_id: str,
98
+ *,
99
+ revision: str,
100
+ backend: str | None = None,
101
+ ) -> AXDocLayout:
102
+ if not isinstance(revision, str) or HUB_COMMIT_RE.fullmatch(revision) is None:
103
+ raise InvalidRevisionError(
104
+ "revision must be a 40-character immutable Hub commit; "
105
+ "'main', a tag, or 'latest' is not accepted"
106
+ )
107
+ if not isinstance(backend, str) or backend not in BACKENDS:
108
+ raise BackendUnavailableError(
109
+ "backend must be given explicitly as one of pytorch, onnx"
110
+ )
111
+ receipt = PINNED_HUB_RECEIPT
112
+ if not isinstance(model_id, str) or model_id != MODEL_ID:
113
+ raise ArtifactError("artifact_model_id_invalid")
114
+ if receipt is None or revision != receipt.model_revision:
115
+ raise ArtifactUnavailableError("no trusted artifact for this revision")
116
+ try:
117
+ dependencies = ("torch", "safetensors") if backend == "pytorch" else ("onnxruntime",)
118
+ for dependency in dependencies:
119
+ try:
120
+ importlib.import_module(dependency)
121
+ except (ImportError, OSError):
122
+ raise BackendUnavailableError("requested backend is not installed") from None
123
+ with TemporaryDirectory(prefix="ax-doclayout-hub-") as temporary:
124
+ directory = Path(temporary)
125
+ _download_bundle(directory, receipt=receipt)
126
+ qualified = qualify_published_bundle(
127
+ model_id=model_id, revision=revision, backend=backend,
128
+ receipt=receipt, bundle=directory,
129
+ )
130
+ if qualified is None:
131
+ raise ArtifactUnavailableError("no trusted artifact for this revision")
132
+ predictor = _load_backend(qualified, backend=backend)
133
+ except AXDocLayoutError:
134
+ raise
135
+ except Exception: # noqa: BLE001 - keep public errors content-free
136
+ raise ArtifactError("hub_backend_load_failed") from None
137
+ model = cls(model_id, revision, backend)
138
+ model._predict_elements = predictor
139
+ model._identity = ModelIdentity(
140
+ id=receipt.model_id, revision=receipt.model_revision,
141
+ artifact_sha256=receipt.artifact_sha256.removeprefix("sha256:"),
142
+ )
143
+ return model
144
+
145
+ def predict(self, page: PageImage) -> PageStructure:
146
+ if self._predict_elements is None or self._identity is None:
147
+ raise ArtifactUnavailableError("no loaded artifact; construct through from_pretrained")
148
+ if not isinstance(page, PageImage):
149
+ raise InvalidPageImageError("page must be a PageImage")
150
+ try:
151
+ elements = self._predict_elements(page)
152
+ return page_structure_from_elements(
153
+ page=page, elements=elements, identity=self._identity,
154
+ )
155
+ except AXDocLayoutError:
156
+ raise
157
+ except Exception: # noqa: BLE001 - keep public errors content-free
158
+ raise InferenceError("page_inference_failed") from None
inference/src/ax_doclayout/postprocess.py ADDED
@@ -0,0 +1,126 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Map letterbox-space detector outputs onto canonical ``PageStructure`` elements.
5
+
6
+ Class id 0 is reserved and never returned. Names that exist only on the
7
+ DocLayout-YOLO incumbent (``figure_caption``, ``table_caption``, and similar)
8
+ are not invented here.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ from collections.abc import Sequence
14
+
15
+ import numpy as np
16
+
17
+ from ax_doclayout.contract import (
18
+ CLASS_NAMES,
19
+ DEFAULT_THRESHOLD_PROFILE,
20
+ BoundingBox,
21
+ Element,
22
+ InvalidPageStructureError,
23
+ ModelIdentity,
24
+ PageImage,
25
+ PageStructure,
26
+ )
27
+ from ax_doclayout.preprocess import Letterbox
28
+
29
+ DEFAULT_SCORE_THRESHOLD = 0.30
30
+
31
+
32
+ def elements_from_letterbox_detections(
33
+ *,
34
+ labels: object,
35
+ boxes: object,
36
+ scores: object,
37
+ page: PageImage,
38
+ letterbox: Letterbox,
39
+ threshold: float = DEFAULT_SCORE_THRESHOLD,
40
+ ) -> tuple[Element, ...]:
41
+ """Keep finite, in-range detections and invert them onto the original raster."""
42
+
43
+ try:
44
+ threshold_value = float(threshold)
45
+ except (TypeError, ValueError) as exc:
46
+ raise InvalidPageStructureError("threshold must be within [0, 1]") from exc
47
+ if not 0.0 <= threshold_value <= 1.0:
48
+ raise InvalidPageStructureError("threshold must be within [0, 1]")
49
+ try:
50
+ label_array = np.asarray(labels)
51
+ box_array = np.asarray(boxes, dtype=np.float64)
52
+ score_array = np.asarray(scores, dtype=np.float64)
53
+ except (TypeError, ValueError) as exc:
54
+ raise InvalidPageStructureError(
55
+ "detection arrays must be rectangular numeric arrays"
56
+ ) from exc
57
+ if label_array.ndim == 0 or box_array.ndim == 0 or score_array.ndim == 0:
58
+ raise InvalidPageStructureError("detection arrays must share a length")
59
+ if label_array.ndim == 2 and label_array.shape[0] == 1:
60
+ if box_array.ndim != 3 or score_array.ndim != 2:
61
+ raise InvalidPageStructureError("detection arrays must share a batch dimension")
62
+ if box_array.shape[0] != 1 or score_array.shape[0] != 1:
63
+ raise InvalidPageStructureError("detection arrays must share a batch dimension")
64
+ label_array = label_array[0]
65
+ box_array = box_array[0]
66
+ score_array = score_array[0]
67
+ if box_array.ndim != 2 or box_array.shape[1] != 4:
68
+ raise InvalidPageStructureError("boxes must be an N x 4 xyxy array")
69
+ if label_array.ndim != 1 or score_array.ndim != 1:
70
+ raise InvalidPageStructureError("detection arrays must share a length")
71
+ if label_array.shape[0] != box_array.shape[0] or label_array.shape[0] != score_array.shape[0]:
72
+ raise InvalidPageStructureError("detection arrays must share a length")
73
+ if not np.isfinite(letterbox.scale) or letterbox.scale <= 0.0:
74
+ return ()
75
+ # Apply identical float64 geometry in bulk before allocating Python objects.
76
+ valid = (
77
+ np.isfinite(box_array).all(axis=1)
78
+ & np.isfinite(score_array)
79
+ & (score_array >= threshold_value)
80
+ & (score_array <= 1.0)
81
+ )
82
+ indices = np.flatnonzero(valid)
83
+ mapped = box_array[indices].copy()
84
+ mapped -= [letterbox.pad_left, letterbox.pad_top,
85
+ letterbox.pad_left, letterbox.pad_top]
86
+ mapped /= letterbox.scale
87
+ np.clip(mapped, 0.0, [page.width_px, page.height_px,
88
+ page.width_px, page.height_px], out=mapped)
89
+ nonempty = (mapped[:, 0] < mapped[:, 2]) & (mapped[:, 1] < mapped[:, 3])
90
+ kept: list[Element] = []
91
+ for index, box in zip(indices[nonempty], mapped[nonempty], strict=True):
92
+ try:
93
+ class_id = int(label_array[index])
94
+ except (TypeError, ValueError, OverflowError):
95
+ continue
96
+ if class_id not in CLASS_NAMES:
97
+ continue
98
+ kept.append(
99
+ Element(
100
+ id=f"{page.page_index}:{len(kept)}",
101
+ class_id=class_id,
102
+ confidence=float(score_array[index]),
103
+ bbox=BoundingBox(*(float(value) for value in box)),
104
+ )
105
+ )
106
+ return tuple(kept)
107
+
108
+
109
+ def page_structure_from_elements(
110
+ *,
111
+ page: PageImage,
112
+ elements: Sequence[Element],
113
+ identity: ModelIdentity,
114
+ threshold_override: float | None = None,
115
+ ) -> PageStructure:
116
+ """Bind already-inverted elements to a caller-supplied model identity."""
117
+
118
+ return PageStructure(
119
+ model=identity,
120
+ page_index=page.page_index,
121
+ width_px=page.width_px,
122
+ height_px=page.height_px,
123
+ elements=tuple(elements),
124
+ threshold_profile=DEFAULT_THRESHOLD_PROFILE,
125
+ threshold_override=threshold_override,
126
+ )
inference/src/ax_doclayout/preprocess.py ADDED
@@ -0,0 +1,145 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Public ``rgb-letterbox-1024/1`` geometry. Independent of any detector graph.
5
+
6
+ This is the signed AXDocLayout-M v1.2 preprocess: HWC uint8 RGB, bilinear resize of the
7
+ longer side to 1024, white pad, odd remainder bottom/right, then inverse
8
+ mapping of letterbox-space boxes back onto the original raster.
9
+
10
+ It is not a DocLayout-YOLO drop-in. That incumbent uses a different colour
11
+ order, a different pad value, and a different class vocabulary. Those
12
+ differences are left unmapped on purpose.
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ from dataclasses import dataclass
18
+
19
+ import numpy as np
20
+ from PIL import Image
21
+
22
+ from ax_doclayout.contract import PREPROCESS_ID, PageImage
23
+ from ax_doclayout.errors import InvalidPageImageError
24
+
25
+ SIZE = 1024
26
+ PAD_RGB = (255, 255, 255)
27
+ FILL_SCALE = 1.0 / 255.0
28
+
29
+
30
+ @dataclass(frozen=True, slots=True)
31
+ class Letterbox:
32
+ """One 1024-square canvas plus the exact inverse-mapping terms."""
33
+
34
+ canvas: np.ndarray
35
+ scale: float
36
+ pad_left: int
37
+ pad_top: int
38
+ pad_right: int
39
+ pad_bottom: int
40
+ size: int = SIZE
41
+
42
+ def __post_init__(self) -> None:
43
+ canvas = np.ascontiguousarray(self.canvas)
44
+ if canvas.ndim != 3 or canvas.shape != (self.size, self.size, 3):
45
+ raise InvalidPageImageError("letterbox canvas must be size x size x 3")
46
+ if canvas.dtype != np.uint8:
47
+ raise InvalidPageImageError("letterbox canvas must be uint8")
48
+ object.__setattr__(self, "canvas", canvas)
49
+
50
+
51
+ def letterbox_rgb(rgb: np.ndarray, *, size: int = SIZE) -> Letterbox:
52
+ """Letterbox a decoded HWC uint8 RGB array under ``rgb-letterbox-1024/1``."""
53
+
54
+ if not isinstance(rgb, np.ndarray):
55
+ raise InvalidPageImageError("rgb must be a numpy array")
56
+ if rgb.ndim != 3 or rgb.shape[2] != 3:
57
+ raise InvalidPageImageError("rgb must be HWC with three channels")
58
+ if rgb.dtype != np.uint8:
59
+ raise InvalidPageImageError("rgb must be uint8")
60
+ height, width = (int(rgb.shape[0]), int(rgb.shape[1]))
61
+ if height < 1 or width < 1:
62
+ raise InvalidPageImageError("rgb must not be empty")
63
+ scale = size / max(width, height)
64
+ new_width = max(1, round(width * scale))
65
+ new_height = max(1, round(height * scale))
66
+ image = Image.fromarray(np.ascontiguousarray(rgb), mode="RGB")
67
+ resized = image.resize((new_width, new_height), Image.BILINEAR)
68
+ pad_width = size - new_width
69
+ pad_height = size - new_height
70
+ pad_left = pad_width // 2
71
+ pad_right = pad_width - pad_left
72
+ pad_top = pad_height // 2
73
+ pad_bottom = pad_height - pad_top
74
+ if pad_right < 0 or pad_bottom < 0:
75
+ raise InvalidPageImageError("letterbox overflow")
76
+ canvas = Image.new("RGB", (size, size), PAD_RGB)
77
+ canvas.paste(resized, (pad_left, pad_top))
78
+ return Letterbox(
79
+ canvas=np.asarray(canvas, dtype=np.uint8),
80
+ scale=float(scale),
81
+ pad_left=int(pad_left),
82
+ pad_top=int(pad_top),
83
+ pad_right=int(pad_right),
84
+ pad_bottom=int(pad_bottom),
85
+ size=int(size),
86
+ )
87
+
88
+
89
+ def letterbox_page(page: PageImage, *, size: int = SIZE) -> Letterbox:
90
+ """Letterbox a validated ``PageImage``."""
91
+
92
+ return letterbox_rgb(page.rgb, size=size)
93
+
94
+
95
+ def letterbox_nchw(letterbox: Letterbox) -> np.ndarray:
96
+ """Return a batch-1 NCHW float32 tensor in ``[0, 1]``."""
97
+
98
+ chw = np.transpose(letterbox.canvas, (2, 0, 1)).astype(np.float32)
99
+ chw *= np.float32(FILL_SCALE)
100
+ return np.ascontiguousarray(chw[None, ...])
101
+
102
+
103
+ def inverse_letterbox_box(
104
+ box: tuple[float, float, float, float] | list[float],
105
+ letterbox: Letterbox,
106
+ *,
107
+ width_px: int,
108
+ height_px: int,
109
+ ) -> tuple[float, float, float, float] | None:
110
+ """Map a letterbox xyxy box onto the original raster, or drop it."""
111
+
112
+ if width_px < 1 or height_px < 1:
113
+ return None
114
+ try:
115
+ x_min, y_min, x_max, y_max = (float(value) for value in box)
116
+ except (TypeError, ValueError):
117
+ return None
118
+ if not all(np.isfinite(value) for value in (x_min, y_min, x_max, y_max)):
119
+ return None
120
+ scale = letterbox.scale
121
+ if scale <= 0.0:
122
+ return None
123
+ mapped = (
124
+ (x_min - letterbox.pad_left) / scale,
125
+ (y_min - letterbox.pad_top) / scale,
126
+ (x_max - letterbox.pad_left) / scale,
127
+ (y_max - letterbox.pad_top) / scale,
128
+ )
129
+ clipped_x_min = min(max(mapped[0], 0.0), float(width_px))
130
+ clipped_y_min = min(max(mapped[1], 0.0), float(height_px))
131
+ clipped_x_max = min(max(mapped[2], 0.0), float(width_px))
132
+ clipped_y_max = min(max(mapped[3], 0.0), float(height_px))
133
+ if clipped_x_min >= clipped_x_max or clipped_y_min >= clipped_y_max:
134
+ return None
135
+ return clipped_x_min, clipped_y_min, clipped_x_max, clipped_y_max
136
+
137
+
138
+ def orig_target_sizes(letterbox: Letterbox) -> np.ndarray:
139
+ """Sizes fed to the exported graph. The graph is bound to the letterbox square."""
140
+
141
+ side = np.int64(letterbox.size)
142
+ return np.ascontiguousarray([[side, side]])
143
+
144
+
145
+ PREPROCESS = PREPROCESS_ID
inference/src/ax_doclayout/py.typed ADDED
File without changes
inference/src/ax_doclayout/taxonomy.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 AutomatosX
2
+ # SPDX-License-Identifier: AGPL-3.0-only
3
+ # Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
4
+ """Canonical 11-class layout taxonomy."""
5
+ from __future__ import annotations
6
+
7
+ from ax_doclayout.errors import TaxonomyError
8
+
9
+ TAXONOMY_ID = "ax-doclayout-11/1"
10
+ RESERVED_ID = 0
11
+
12
+ CLASS_NAMES: dict[int, str] = {
13
+ 1: "caption",
14
+ 2: "footnote",
15
+ 3: "formula",
16
+ 4: "list_item",
17
+ 5: "page_footer",
18
+ 6: "page_header",
19
+ 7: "picture",
20
+ 8: "section_header",
21
+ 9: "table",
22
+ 10: "text",
23
+ 11: "title",
24
+ }
25
+
26
+ NAME_TO_ID: dict[str, int] = {name: class_id for class_id, name in CLASS_NAMES.items()}
27
+
28
+
29
+ def public_name(class_id: int) -> str:
30
+ if class_id == RESERVED_ID:
31
+ raise TaxonomyError("class_id_reserved")
32
+ try:
33
+ return CLASS_NAMES[class_id]
34
+ except KeyError as exc:
35
+ raise TaxonomyError("class_id_unknown") from exc
36
+
37
+
38
+ def public_id(class_name: str) -> int:
39
+ try:
40
+ return NAME_TO_ID[class_name]
41
+ except KeyError as exc:
42
+ raise TaxonomyError("class_name_unknown") from exc
43
+
44
+
45
+ def to_internal(class_id: int) -> int:
46
+ """Map public 1–11 to zero-based indices 0–10."""
47
+ if class_id not in CLASS_NAMES:
48
+ raise TaxonomyError("class_id_unknown")
49
+ return class_id - 1
50
+
51
+
52
+ def from_internal(index: int) -> int:
53
+ class_id = index + 1
54
+ if class_id not in CLASS_NAMES:
55
+ raise TaxonomyError("class_id_unknown")
56
+ return class_id