https://huggingface.co/nvidia/nemotron-ocr-v1/tree/main

LICENSE

@@ -1,243 +0,0 @@
-All binary model files are licensed under NVIDIA Open Model License Agreement.
-All source code files are licensed under the Apache 2.0 License. 
-
-------------
-NVIDIA Open Model License Agreement
-Last Modified: October 24, 2025
-https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/
-
-This NVIDIA Open Model License Agreement (the “Agreement”) is a legal agreement between the Legal Entity You represent, or if no entity is identified, You and NVIDIA Corporation and its Affiliates (“NVIDIA”) and governs Your use of the Models that NVIDIA provides to You under this Agreement. NVIDIA and You are each a “party” and collectively the “parties.”
-
-NVIDIA models released under this Agreement are intended to be used permissively and enable the further development of AI technologies. Subject to the terms of this Agreement, NVIDIA confirms that:
-
-- Models are commercially usable.
-- You are free to create and distribute Derivative Models.
-- NVIDIA does not claim ownership to any outputs generated using the Models or Derivative Models.
-By using, reproducing, modifying, distributing, performing or displaying any portion or element of the Model or Derivative Model, or otherwise accepting the terms of this Agreement, you agree to be bound by this Agreement.
-
-1. Definitions. The following definitions apply to this Agreement:
-
-1.1 "Derivative Model" means all (a) modifications to the Model, (b) works based on the Model, and (c) any other derivative works of the Model. An output is not a Derivative Model.
-
-1.2 "Legal Entity" means the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (a) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (b) ownership of fifty percent (50%) or more of the outstanding shares, or (c) beneficial ownership of such entity.
-
-1.3 “Model” means the machine learning model, software, checkpoints, learnt weights, algorithms, parameters, configuration files and documentation shared under this Agreement.
-
-1.4 "NVIDIA Cosmos Model" means a multimodal Model shared under this Agreement
-
-1.5 "Special-Purpose Model" means a Model that is only competent in a narrow set of purpose-specific tasks and should not be used for unintended or general-purpose applications
-
-1.6 “You” or “Your” means an individual or Legal Entity exercising permissions granted by this Agreement.
-
-2. Conditions for Use, License Grant, AI Ethics and IP Ownership.
-
-2.1 Conditions for Use. The Model and any Derivative Model are subject to additional terms as described in Section 2 and Section 3 of this Agreement and govern Your use. If You institute copyright or patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Model or a Derivative Model constitutes direct or contributory copyright or patent infringement, then any licenses granted to You under this Agreement for that Model or Derivative Model will terminate as of the date such litigation is filed. If You bypass, disable, reduce the efficacy of, or circumvent any technical limitation, safety guardrail or associated safety guardrail hyperparameter, encryption, security, digital rights management, or authentication mechanism (collectively “Guardrail”) contained in the Model without a substantially similar Guardrail appropriate for your use case, your rights under this Agreement will automatically terminate. NVIDIA may indicate in relevant documentation that a Model is a Special-Purpose Model. NVIDIA may update this Agreement to comply with legal and regulatory requirements at any time and You agree to either comply with any updated license or cease Your copying, use, and distribution of the Model and any Derivative Model.
-
-2.2 License Grant. The rights granted herein are explicitly conditioned on Your full compliance with the terms of this Agreement. Subject to the terms and conditions of this Agreement, NVIDIA hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, revocable (as stated in Section 2.1) license to publicly perform, publicly display, reproduce, use, create derivative works of, make, have made, sell, offer for sale, distribute (through multiple tiers of distribution) and import the Model.
-
-2.3 AI Ethics. Use of the Models under the Agreement must be consistent with NVIDIA’s Trustworthy AI terms found at https://www.nvidia.com/en-us/agreements/trustworthy-ai/terms/.
-
-2.4 NVIDIA owns the Model and any Derivative Models created by NVIDIA. Subject to NVIDIA’s underlying ownership rights in the Model or its Derivative Models, You are and will be the owner of Your Derivative Models. NVIDIA claims no ownership rights in outputs. You are responsible for outputs and their subsequent uses. Except as expressly granted in this Agreement, (a) NVIDIA reserves all rights, interests and remedies in connection with the Model and (b) no other license or right is granted to you by implication, estoppel or otherwise.
-
-3. Redistribution. You may reproduce and distribute copies of the Model or Derivative Models thereof in any medium, with or without modifications, provided that You meet the following conditions:
-
-3.1 If you distribute the Model, You must give any other recipients of the Model a copy of this Agreement and include the following attribution notice within a “Notice” text file with such copies: “Licensed by NVIDIA Corporation under the NVIDIA Open Model License”;
-
-3.2 If you distribute or make available a NVIDIA Cosmos Model, or a product or service (including an AI model) that contains or uses a NVIDIA Cosmos Model, use a NVIDIA Cosmos Model to create a Derivative Model, or use a NVIDIA Cosmos Model or its outputs to create, train, fine tune, or otherwise improve an AI model, you will include “Built on NVIDIA Cosmos” on a related website, user interface, blogpost, about page, or product documentation; and
-
-3.3 You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Models as a whole, provided Your use, reproduction, and distribution of the Model otherwise complies with the conditions stated in this Agreement.
-
-4. Separate Components. The Models may include or be distributed with components provided with separate legal notices or terms that accompany the components, such as an Open Source Software License or other third-party license. The components are subject to the applicable other licenses, including any proprietary notices, disclaimers, requirements and extended use rights; except that this Agreement will prevail regarding the use of third-party Open Source Software License, unless a third-party Open Source Software License requires its license terms to prevail. “Open Source Software License” means any software, data or documentation subject to any license identified as an open source license by the Open Source Initiative (https://opensource.org), Free Software Foundation (https://www.fsf.org) or other similar open source organization or listed by the Software Package Data Exchange (SPDX) Workgroup under the Linux Foundation (https://www.spdx.org).
-
-5. Trademarks. This Agreement does not grant permission to use the trade names, trademarks, service marks, or product names of NVIDIA, except as required for reasonable and customary use in describing the origin of the Model and reproducing the content of the “Notice” text file.
-
-6. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, NVIDIA provides the Model on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for reviewing Model documentation, including any Special-Purpose Model limitations, and determining the appropriateness of using or redistributing the Model, Derivative Models and outputs. You assume any risks associated with Your exercise of permissions under this Agreement.
-
-7. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, will NVIDIA be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this Agreement or out of the use or inability to use the Model, Derivative Models or outputs (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if NVIDIA has been advised of the possibility of such damages.
-
-8. Indemnity. You will indemnify and hold harmless NVIDIA from and against any claim by any third party arising out of or related to your use or distribution of the Model, Derivative Models or outputs.
-
-9. Feedback. NVIDIA appreciates your feedback, and You agree that NVIDIA may use it without restriction or compensation to You.
-
-10. Governing Law. This Agreement will be governed in all respects by the laws of the United States and the laws of the State of Delaware, without regard to conflict of laws principles or the United Nations Convention on Contracts for the International Sale of Goods. The state and federal courts residing in Santa Clara County, California will have exclusive jurisdiction over any dispute or claim arising out of or related to this Agreement, and the parties irrevocably consent to personal jurisdiction and venue in those courts; except that, either party may apply for injunctive remedies or an equivalent type of urgent legal relief in any jurisdiction.
-
-11. Trade and Compliance. You agree to comply with all applicable export, import, trade and economic sanctions laws and regulations, as amended, including without limitation U.S. Export Administration Regulations and Office of Foreign Assets Control regulations. These laws include restrictions on destinations, end-users and end-use.
-
-Version Release Date: October 24, 2025
-
------------------
-                                 Apache License
-                           Version 2.0, January 2004
-                        http://www.apache.org/licenses/
-   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
-   1. Definitions.
-      "License" shall mean the terms and conditions for use, reproduction,
-      and distribution as defined by Sections 1 through 9 of this document.
-      "Licensor" shall mean the copyright owner or entity authorized by
-      the copyright owner that is granting the License.
-      "Legal Entity" shall mean the union of the acting entity and all
-      other entities that control, are controlled by, or are under common
-      control with that entity. For the purposes of this definition,
-      "control" means (i) the power, direct or indirect, to cause the
-      direction or management of such entity, whether by contract or
-      otherwise, or (ii) ownership of fifty percent (50%) or more of the
-      outstanding shares, or (iii) beneficial ownership of such entity.
-      "You" (or "Your") shall mean an individual or Legal Entity
-      exercising permissions granted by this License.
-      "Source" form shall mean the preferred form for making modifications,
-      including but not limited to software source code, documentation
-      source, and configuration files.
-      "Object" form shall mean any form resulting from mechanical
-      transformation or translation of a Source form, including but
-      not limited to compiled object code, generated documentation,
-      and conversions to other media types.
-      "Work" shall mean the work of authorship, whether in Source or
-      Object form, made available under the License, as indicated by a
-      copyright notice that is included in or attached to the work
-      (an example is provided in the Appendix below).
-      "Derivative Works" shall mean any work, whether in Source or Object
-      form, that is based on (or derived from) the Work and for which the
-      editorial revisions, annotations, elaborations, or other modifications
-      represent, as a whole, an original work of authorship. For the purposes
-      of this License, Derivative Works shall not include works that remain
-      separable from, or merely link (or bind by name) to the interfaces of,
-      the Work and Derivative Works thereof.
-      "Contribution" shall mean any work of authorship, including
-      the original version of the Work and any modifications or additions
-      to that Work or Derivative Works thereof, that is intentionally
-      submitted to Licensor for inclusion in the Work by the copyright owner
-      or by an individual or Legal Entity authorized to submit on behalf of
-      the copyright owner. For the purposes of this definition, "submitted"
-      means any form of electronic, verbal, or written communication sent
-      to the Licensor or its representatives, including but not limited to
-      communication on electronic mailing lists, source code control systems,
-      and issue tracking systems that are managed by, or on behalf of, the
-      Licensor for the purpose of discussing and improving the Work, but
-      excluding communication that is conspicuously marked or otherwise
-      designated in writing by the copyright owner as "Not a Contribution."
-      "Contributor" shall mean Licensor and any individual or Legal Entity
-      on behalf of whom a Contribution has been received by Licensor and
-      subsequently incorporated within the Work.
-   2. Grant of Copyright License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      copyright license to reproduce, prepare Derivative Works of,
-      publicly display, publicly perform, sublicense, and distribute the
-      Work and such Derivative Works in Source or Object form.
-   3. Grant of Patent License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      (except as stated in this section) patent license to make, have made,
-      use, offer to sell, sell, import, and otherwise transfer the Work,
-      where such license applies only to those patent claims licensable
-      by such Contributor that are necessarily infringed by their
-      Contribution(s) alone or by combination of their Contribution(s)
-      with the Work to which such Contribution(s) was submitted. If You
-      institute patent litigation against any entity (including a
-      cross-claim or counterclaim in a lawsuit) alleging that the Work
-      or a Contribution incorporated within the Work constitutes direct
-      or contributory patent infringement, then any patent licenses
-      granted to You under this License for that Work shall terminate
-      as of the date such litigation is filed.
-   4. Redistribution. You may reproduce and distribute copies of the
-      Work or Derivative Works thereof in any medium, with or without
-      modifications, and in Source or Object form, provided that You
-      meet the following conditions:
-      (a) You must give any other recipients of the Work or
-          Derivative Works a copy of this License; and
-      (b) You must cause any modified files to carry prominent notices
-          stating that You changed the files; and
-      (c) You must retain, in the Source form of any Derivative Works
-          that You distribute, all copyright, patent, trademark, and
-          attribution notices from the Source form of the Work,
-          excluding those notices that do not pertain to any part of
-          the Derivative Works; and
-      (d) If the Work includes a "NOTICE" text file as part of its
-          distribution, then any Derivative Works that You distribute must
-          include a readable copy of the attribution notices contained
-          within such NOTICE file, excluding those notices that do not
-          pertain to any part of the Derivative Works, in at least one
-          of the following places: within a NOTICE text file distributed
-          as part of the Derivative Works; within the Source form or
-          documentation, if provided along with the Derivative Works; or,
-          within a display generated by the Derivative Works, if and
-          wherever such third-party notices normally appear. The contents
-          of the NOTICE file are for informational purposes only and
-          do not modify the License. You may add Your own attribution
-          notices within Derivative Works that You distribute, alongside
-          or as an addendum to the NOTICE text from the Work, provided
-          that such additional attribution notices cannot be construed
-          as modifying the License.
-      You may add Your own copyright statement to Your modifications and
-      may provide additional or different license terms and conditions
-      for use, reproduction, or distribution of Your modifications, or
-      for any such Derivative Works as a whole, provided Your use,
-      reproduction, and distribution of the Work otherwise complies with
-      the conditions stated in this License.
-   5. Submission of Contributions. Unless You explicitly state otherwise,
-      any Contribution intentionally submitted for inclusion in the Work
-      by You to the Licensor shall be under the terms and conditions of
-      this License, without any additional terms or conditions.
-      Notwithstanding the above, nothing herein shall supersede or modify
-      the terms of any separate license agreement you may have executed
-      with Licensor regarding such Contributions.
-   6. Trademarks. This License does not grant permission to use the trade
-      names, trademarks, service marks, or product names of the Licensor,
-      except as required for reasonable and customary use in describing the
-      origin of the Work and reproducing the content of the NOTICE file.
-   7. Disclaimer of Warranty. Unless required by applicable law or
-      agreed to in writing, Licensor provides the Work (and each
-      Contributor provides its Contributions) on an "AS IS" BASIS,
-      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-      implied, including, without limitation, any warranties or conditions
-      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-      PARTICULAR PURPOSE. You are solely responsible for determining the
-      appropriateness of using or redistributing the Work and assume any
-      risks associated with Your exercise of permissions under this License.
-   8. Limitation of Liability. In no event and under no legal theory,
-      whether in tort (including negligence), contract, or otherwise,
-      unless required by applicable law (such as deliberate and grossly
-      negligent acts) or agreed to in writing, shall any Contributor be
-      liable to You for damages, including any direct, indirect, special,
-      incidental, or consequential damages of any character arising as a
-      result of this License or out of the use or inability to use the
-      Work (including but not limited to damages for loss of goodwill,
-      work stoppage, computer failure or malfunction, or any and all
-      other commercial damages or losses), even if such Contributor
-      has been advised of the possibility of such damages.
-   9. Accepting Warranty or Additional Liability. While redistributing
-      the Work or Derivative Works thereof, You may choose to offer,
-      and charge a fee for, acceptance of support, warranty, indemnity,
-      or other liability obligations and/or rights consistent with this
-      License. However, in accepting such obligations, You may act only
-      on Your own behalf and on Your sole responsibility, not on behalf
-      of any other Contributor, and only if You agree to indemnify,
-      defend, and hold each Contributor harmless for any liability
-      incurred by, or claims asserted against, such Contributor by reason
-      of your accepting any such warranty or additional liability.
-   END OF TERMS AND CONDITIONS
-   APPENDIX: How to apply the Apache License to your work.
-      To apply the Apache License to your work, attach the following
-      boilerplate notice, with the fields enclosed by brackets "[]"
-      replaced with your own identifying information. (Don't include
-      the brackets!)  The text should be enclosed in the appropriate
-      comment syntax for the file format. We also recommend that a
-      file or class name and description of purpose be included on the
-      same "printed page" as the copyright notice for easier
-      identification within third-party archives.
-
-   Copyright [yyyy] [name of copyright owner]
-
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
-
-       http://www.apache.org/licenses/LICENSE-2.0
-
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License.
-
-   

README.md

@@ -14,7 +14,6 @@ tags:
 - text recognition
 - layout analysis
 - ingestion
-new_version: nvidia/nemotron-ocr-v2
 ---
 
 # Nemotron OCR v1
@@ -346,4 +345,4 @@ Please report security vulnerabilities or NVIDIA AI Concerns [here](https://app.
 | Model Application Field(s): | Text recognition and structured OCR for multimodal retrieval. Inputs can include natural scene images, scanned documents, charts, tables, and infographics. |
 | Use Case Restrictions: | Abide by [NVIDIA Open Model License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/) and the use of the post-processing scripts are licensed under [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0.txt). |
 | Model and dataset restrictions: | The principle of least privilege (PoLP) is applied, limiting access for dataset generation and model development. Restrictions enforce dataset access only during training, and all dataset license constraints are adhered to. |
-| Describe the life critical impact (if present): | Not applicable. |
+| Describe the life critical impact (if present): | Not applicable. |

config.json

@@ -1 +0,0 @@
-{}

example.py

@@ -8,7 +8,7 @@ from nemotron_ocr.inference.pipeline import NemotronOCR
 
 
 def main(image_path, merge_level, no_visualize, model_dir):
-    ocr_pipeline = NemotronOCR(model_dir=model_dir)
+    ocr_pipeline = NemotronOCR()
 
     predictions = ocr_pipeline(image_path, merge_level=merge_level, visualize=not no_visualize)
 

nemotron-ocr/cpp/non_maximal_suppression/cuda_non_maximal_suppression.cu

@@ -157,8 +157,11 @@ void device_row_collapse(torch::PackedTensorAccessor64<T, 5> allQuads,
     torch::PackedTensorAccessor64<T, 3> allConfs,
     T confThreshold, T iouThreshold,
     torch::PackedTensorAccessor64<int32_t, 1> allOutCounts,
-    torch::PackedTensorAccessor64<T, 3> allOutEmbedQuads,
-    torch::PackedTensorAccessor64<int32_t, 2> allOutIds)
+    torch::PackedTensorAccessor64<T, 3> allOutEmbedQuads
+#ifdef NMS_VERIFY_CORRECTNESS
+    , torch::PackedTensorAccessor64<int32_t, 2> allOutIds
+#endif
+    )
 {
     typedef InPlaceQuad_<T> Quadf;
     static_assert(sizeof(Quadf) == sizeof(T) * 8, "Invalid QuadMem size!");
@@ -303,9 +306,11 @@ void device_row_collapse(torch::PackedTensorAccessor64<T, 5> allQuads,
         }
 
         write_embed_quad(outEmbedQuads, outQuad, storeOff + procLabel - 1);
+#ifdef NMS_VERIFY_CORRECTNESS
         if (threadRank == 0) {
             allOutIds[b][storeOff + procLabel - 1] = r * 32 + startIdx;
         }
+#endif
     }
 
     if (threadRank == 0) {
@@ -316,9 +321,9 @@ void device_row_collapse(torch::PackedTensorAccessor64<T, 5> allQuads,
 #undef threadRank
 }
 
-template<typename T>
+template<bool IsSingleExample, typename T>
 __global__
-void device_a2a_adjacency_sparse(const int32_t *ptrQuadCts,
+void device_a2a_adjacency_sparse(const uint64_t punCounts,
                                  T iouThreshold,
                                  torch::PackedTensorAccessor64<T, 3> embedQuads,
                                  torch::PackedTensorAccessor64<bool, 2> outIsStart,
@@ -327,11 +332,7 @@ void device_a2a_adjacency_sparse(const int32_t *ptrQuadCts,
 {
     const uint32_t b = blockIdx.y;
 
-    const int32_t quadCt = ptrQuadCts[b];
-
-    if (quadCt == 0) {
-        return;
-    }
+    const int32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
 
     const int32_t jobIdx = blockIdx.x * blockDim.x + threadIdx.x;
     const int32_t row = jobIdx / quadCt;
@@ -342,7 +343,7 @@ void device_a2a_adjacency_sparse(const int32_t *ptrQuadCts,
         return;
     }
 
-    T* exData = embedQuads[b].data();
+    T* exData = IsSingleExample ? embedQuads.data() : embedQuads[b].data();
 
     const auto qRow = StridedEmbedQuad_<T>{ exData + row * embedQuads.stride(2), embedQuads.stride(1) }.Bounds(),
                qCol = StridedEmbedQuad_<T>{ exData + col * embedQuads.stride(2), embedQuads.stride(1) }.Bounds();
@@ -404,9 +405,9 @@ void device_a2a_adjacency_sparse(const int32_t *ptrQuadCts,
     }
 }
 
-template<uint32_t NumWarps, typename T, int32_t I_CELL_SIZE>
+template<uint32_t NumWarps, bool IsSingleExample, typename T, int32_t I_CELL_SIZE>
 __global__
-void device_a2a_adjacency_build_grid(const int32_t *ptrQuadCts,
+void device_a2a_adjacency_build_grid(const uint64_t punCounts,
                                      torch::PackedTensorAccessor64<T, 3> embedQuads,
                                      torch::PackedTensorAccessor64<int32_t, 4> outGridCells,
                                      torch::PackedTensorAccessor64<int32_t, 3> outQuadCells)
@@ -422,10 +423,10 @@ void device_a2a_adjacency_build_grid(const int32_t *ptrQuadCts,
 
     const uint32_t b = blockIdx.z;
 
-    const uint32_t quadCt = ptrQuadCts[b];
+    const uint32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
     const uint32_t quadIdx = blockIdx.y;
 
-    if (quadIdx >= quadCt) {
+    if (!IsSingleExample && quadIdx >= quadCt) {
         return;
     }
 
@@ -484,9 +485,9 @@ void device_a2a_adjacency_build_grid(const int32_t *ptrQuadCts,
 
 typedef uint8_t visit_mask_t;
 
-template<uint32_t NumWarps, typename T>
+template<uint32_t NumWarps, bool IsSingleExample, typename T>
 __global__
-void device_a2a_adjacency_with_grid(const int32_t *ptrQuadCts,
+void device_a2a_adjacency_with_grid(const uint64_t punCounts,
                                     T iouThreshold,
                                     torch::PackedTensorAccessor64<T, 3> allEmbedQuads,
                                     torch::PackedTensorAccessor64<int32_t, 4> allCells,
@@ -502,10 +503,10 @@ void device_a2a_adjacency_with_grid(const int32_t *ptrQuadCts,
 
     const uint32_t b = blockIdx.z;
 
-    const uint32_t quadCt = ptrQuadCts[b];
+    const uint32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
     const uint32_t quadIdx = blockIdx.y;
 
-    if (quadIdx >= quadCt) {
+    if (!IsSingleExample && quadIdx >= quadCt) {
         return;
     }
 
@@ -534,7 +535,7 @@ void device_a2a_adjacency_with_grid(const int32_t *ptrQuadCts,
     auto exAdjCounts = reinterpret_cast<uint32_t*>(outAdjCounts[b].data());
     auto exAdjValues = outSparseAdj[b][quadIdx].data();
 
-    T *exData = allEmbedQuads[b].data();
+    T *exData = IsSingleExample ? allEmbedQuads.data() : allEmbedQuads[b].data();
 
     const auto bdsAnchor = Quad_<T>{ s_quadVerts }.Bounds();
 
@@ -598,8 +599,9 @@ void device_a2a_adjacency_with_grid(const int32_t *ptrQuadCts,
     }
 }
 
+template<bool IsSingleExample>
 __global__
-void device_flatten_graph_iterative(const int32_t *ptrQuadCts,
+void device_flatten_graph_iterative(const uint64_t punCounts,
                                     torch::PackedTensorAccessor64<bool, 2> allIsStart,
                                     volatile uint32_t *allAdjCounts,
                                     volatile uint32_t *allAdjValues
@@ -620,12 +622,14 @@ void device_flatten_graph_iterative(const int32_t *ptrQuadCts,
     const uint32_t b = blockIdx.z;
     const uint32_t anchorRow = blockIdx.y;
 
-    const uint32_t quadCt = ptrQuadCts[b];
+    const uint32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
 
     // Only need to check this if there are multiple examples, since in the case of a single example,
     // the grid is precisely sized to that quadCt
-    if (anchorRow >= quadCt) {
-        return;
+    if constexpr (!IsSingleExample) {
+        if (anchorRow >= quadCt) {
+            return;
+        }
     }
 
     auto isStart = allIsStart[b].data();
@@ -686,13 +690,12 @@ void device_flatten_graph_iterative(const int32_t *ptrQuadCts,
     visitStack[1] = anchorRow;
 #ifndef NDEBUG
     for (uint32_t i = 2; i < VISIT_STACK_SIZE; ++i) {
-        visitStack[i] = TERM_VALUE;
+        visitStack[i] = -2;
     }
 #endif
     int32_t visitPtr = 1;
 
-    // NOTE: This loop is actually terminated by the `if (warpNextCol == TERM_VALUE)` check below
-    for (uint32_t dfsIter = 0; true; ++dfsIter) {
+    while (true) {
 #ifdef NMS_VERIFY_CORRECTNESS
         assert(visitPtr >= 0 && visitPtr < VISIT_STACK_SIZE);
 #endif
@@ -704,7 +707,7 @@ void device_flatten_graph_iterative(const int32_t *ptrQuadCts,
         if (threadNextCol == warpNextCol) {
 #ifndef NDEBUG
             // This makes it easier to debug where the pointer is
-            visitStack[visitPtr] = TERM_VALUE;
+            visitStack[visitPtr] = -2;
 #endif
             --visitPtr;
         }
@@ -728,15 +731,12 @@ void device_flatten_graph_iterative(const int32_t *ptrQuadCts,
         const uint32_t procAdjCount = adjCounts[procRow];
         auto procAdjValues = adjValues + (procRow * maxExCount);
 
+        // Offsetting by the iteration number will help balance out the maximum depth of any stack in the warp.
+        // The reason behind this is due to how otherwise, warp-0 will always get a new element, warp-1 iff the adj graph
+        // has more than one element, warp-2 iff the adj graph has more than two elements, and so on. Basically,
+        // the warps have decreasing pressure. With the rotation mechanism, it helps to balance out stack usage.
         for (uint32_t i = threadRank; i < procAdjCount; i += WARP_SIZE) {
-            uint32_t adjCol = procAdjValues[i];
-
-            auto group = cg::coalesced_threads();
-            // Offsetting by the iteration number will help balance out the maximum depth of any stack in the warp.
-            // The reason behind this is due to how otherwise, warp-0 will always get a new element, warp-1 iff the adj graph
-            // has more than one element, warp-2 iff the adj graph has more than two elements, and so on. Basically,
-            // the warps have decreasing pressure. With the rotation mechanism, it helps to balance out stack usage.
-            adjCol = group.shfl(adjCol, (group.thread_rank() + dfsIter) % group.size());
+            const uint32_t adjCol = procAdjValues[i];
 
             // This will set the queued flag for this column, if it's not already set.
             // It also returns the old state. In our case, we only want to add this value to the
@@ -748,6 +748,7 @@ void device_flatten_graph_iterative(const int32_t *ptrQuadCts,
 
             bool alreadyAdded = oldMask & ADDED_MASK;
 
+            auto group = cg::coalesced_threads();
             const uint32_t gThreadRank = group.thread_rank();
             uint32_t notAddedBallot = group.ballot(!alreadyAdded);
             if (notAddedBallot) {
@@ -824,7 +825,8 @@ void add_to_set(const torch::TensorAccessor<int32_t, 1>& adjCounts,
     }
 }
 
-void cpu_flatten_graph(const int32_t *ptrQuadCts,
+template<bool IsSingleExample>
+void cpu_flatten_graph(const uint64_t punCounts,
                        torch::Tensor isStartTensorGPU,
                        torch::Tensor adjCountsTensorGPU,
                        torch::Tensor adjValuesTensorGPU)
@@ -838,7 +840,7 @@ void cpu_flatten_graph(const int32_t *ptrQuadCts,
     auto allAdjValues = adjValuesTensor.accessor<int32_t, 3>();
 
     for (int32_t b = 0; b < allAdjCounts.size(0); ++b) {
-        const int32_t quadCt = ptrQuadCts[b];
+        const int32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
 
         for (int32_t row = 0; row < quadCt; ++row) {
             std::unordered_set<int32_t> fullAdjSet;
@@ -893,9 +895,9 @@ void device_a2a_adj_cleanup(const int32_t *counts,
     }
 }
 
-template<uint32_t NumWarps, typename T>
+template<uint32_t NumWarps, typename T, bool IsSingleExample>
 __global__
-void device_a2a_collapse(torch::PackedTensorAccessor64<int32_t, 1> quadCounts,
+void device_a2a_collapse(const uint64_t punCounts,
                          torch::PackedTensorAccessor64<T, 3> allEmbedQuads,
                          torch::PackedTensorAccessor64<bool, 2> allIsLeadRow,
                          const int64_t *regionCounts,
@@ -915,14 +917,16 @@ void device_a2a_collapse(torch::PackedTensorAccessor64<int32_t, 1> quadCounts,
     const uint32_t b = blockIdx.z;
     const uint32_t row = blockIdx.y;
 
-    const int32_t quadCt = quadCounts[b];
+    const int32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
 
-    if (row >= quadCt) {
-        return;
+    if constexpr (!IsSingleExample) {
+        if (row >= quadCt) {
+            return;
+        }
     }
 
     // Only process the lead rows
-    const auto isLeadRow = allIsLeadRow[b].data();
+    const auto isLeadRow = IsSingleExample ? allIsLeadRow.data() : allIsLeadRow[b].data();
     if (!isLeadRow[row]) {
         return;
     }
@@ -941,7 +945,7 @@ void device_a2a_collapse(torch::PackedTensorAccessor64<int32_t, 1> quadCounts,
         __syncthreads();
     }
 
-    T *exData = allEmbedQuads[b].data();
+    T *exData = IsSingleExample ? allEmbedQuads.data() : allEmbedQuads[b].data();
 
     const int32_t adjCount = allAdjCounts[b][row];
     const int32_t *adjIdxs = allAdjValues[b][row].data();
@@ -982,12 +986,20 @@ void device_a2a_collapse(torch::PackedTensorAccessor64<int32_t, 1> quadCounts,
 
     // Figure out the output position
     uint32_t writePosition = 0;
-    for (int32_t i = threadRank; i < b; i += BLOCK_WIDTH) {
-        writePosition += regionCounts[i];
+    if constexpr (!IsSingleExample) {
+        for (int32_t i = threadRank; i < b; i += BLOCK_WIDTH) {
+            writePosition += regionCounts[i];
+        }
     }
 
+    const int32_t numLongs = row >> 3; // Divide by 8
     const uint8_t *pCurrIsLeadRow = reinterpret_cast<const uint8_t*>(isLeadRow);
-    for (int32_t i = threadRank; i < row; i += BLOCK_WIDTH) {
+    const uint64_t *lpCurrIsLeadRow = reinterpret_cast<const uint64_t*>(pCurrIsLeadRow);
+
+    for (int32_t i = threadRank; i < numLongs; i += BLOCK_WIDTH) {
+        writePosition += __popcll(lpCurrIsLeadRow[i]);
+    }
+    for (int32_t i = (numLongs * 8) + threadRank; i < row; i += BLOCK_WIDTH) {
         if (pCurrIsLeadRow[i]) {
             ++writePosition;
         }
@@ -1063,9 +1075,13 @@ CollapseRowsResult collapse_rows(
     int64_t embedSize = sizeof(EmbedQuad_<scalar_t>) / sizeof(scalar_t);
     auto rowMergeTensor = torch::empty({ quads.size(0), embedSize, quads.size(1) * quads.size(2) }, quads.options());
 
+#ifdef NMS_VERIFY_CORRECTNESS
     auto idsTensor = torch::full({ quads.size(0), quads.size(1) * quads.size(2) },
                                  std::numeric_limits<int32_t>::max(),
                                  counts.options().dtype(torch::kInt32));
+#else
+    torch::Tensor idsTensor;
+#endif
 
     dim3 blockSize(32, 3, 1);
     dim3 gridSize(1,
@@ -1077,8 +1093,10 @@ CollapseRowsResult collapse_rows(
         probs.packed_accessor64<scalar_t, 3>(),
         probThreshold, iouThreshold,
         counts.packed_accessor64<int32_t, 1>(),
-        rowMergeTensor.packed_accessor64<scalar_t, 3>(),
-        idsTensor.packed_accessor64<int32_t, 2>()
+        rowMergeTensor.packed_accessor64<scalar_t, 3>()
+#ifdef NMS_VERIFY_CORRECTNESS
+        , idsTensor.packed_accessor64<int32_t, 2>()
+#endif
     );
 
 #ifdef NMS_VERIFY_CORRECTNESS
@@ -1101,6 +1119,7 @@ CollapseRowsResult collapse_rows(
 
     counts = counts.slice(/*dim=*/ 0, 0, counts.size(0) - 1);
 
+#ifdef NMS_VERIFY_CORRECTNESS
     int64_t maxExCount;
     if (counts.size(0) > 1) {
         maxExCount = counts.max().item<int32_t>();
@@ -1112,12 +1131,13 @@ CollapseRowsResult collapse_rows(
 
     rowMergeTensor = rowMergeTensor.slice(2, 0, maxExCount);
     idsTensor = idsTensor.slice(1, 0, maxExCount);
-    auto order = torch::argsort(idsTensor, /*dim=*/ 1, s_sortOrder);
+    auto order = torch::argsort(idsTensor, /*dim=*/ 1, s_sortOrder); s_sortOrder = !s_sortOrder;
 
     auto embOrder = order.unsqueeze(1).expand_as(rowMergeTensor);
 
     rowMergeTensor = torch::gather(rowMergeTensor, /*dim=*/ 2, embOrder);
     idsTensor = torch::gather(idsTensor, /*dim=*/ 1, order);
+#endif
 
     return { counts, rowMergeTensor, totalQuads, idsTensor, imageWidth, imageHeight };
 }
@@ -1157,8 +1177,8 @@ struct AdjacencyResult {
     int64_t MaxExCount;
 };
 
-template<typename T>
-void cpu_a2a_adjacency_sparse(const int32_t *ptrQuadCts,
+template<bool IsSingleExample, typename T>
+void cpu_a2a_adjacency_sparse(const uint64_t punCounts,
                               const T iouThreshold,
                               torch::Tensor embedQuadsTensor,
                               torch::Tensor outIsStartTensorGPU,
@@ -1176,7 +1196,7 @@ void cpu_a2a_adjacency_sparse(const int32_t *ptrQuadCts,
     auto adjValues = outSparseAdjTensor.accessor<int32_t, 3>();
 
     for (int32_t b = 0; b < embedQuadsTensor.size(0); ++b) {
-        const int32_t quadCt = ptrQuadCts[b];
+        const int32_t quadCt = IsSingleExample ? punCounts : reinterpret_cast<const int32_t*>(punCounts)[b];
 
         T *exData = embedQuads[b].data();
 
@@ -1264,6 +1284,13 @@ AdjacencyResult compute_all_to_all_adjacency(
                                        counts.options().dtype(torch::kInt32));
 #endif
 
+    // If the batch is only a single example, instead of hitting global memory for the count, we can
+    // just encode the count into the pointer instead
+    uint64_t ptrCounts = reinterpret_cast<uint64_t>(counts.data_ptr<int32_t>());
+    if (counts.size(0) == 1) {
+        ptrCounts = maxExCount;
+    }
+
 #ifdef NMS_VERIFY_CORRECTNESS
     auto cpuAdjValuesTensor = adjValuesTensor.cpu();
     auto cpuAdjCountsTensor = adjCountsTensor.cpu();
@@ -1291,15 +1318,23 @@ AdjacencyResult compute_all_to_all_adjacency(
     //blockSize = dim3{ GRID_NUM_WARPS * 32, 1, 1 };
     //gridSize = dim3{ 1, static_cast<uint32_t>(maxExCount), static_cast<uint32_t>(counts.size(0)) };
 
-    //device_a2a_adjacency_build_grid<GRID_NUM_WARPS, scalar_t, CELL_SIZE> KERNEL_ARG2(gridSize, blockSize) (
-    //    counts.data_ptr<int32_t>(),
+    //auto buildGridFn = counts.size(0) == 1 ?
+    //    device_a2a_adjacency_build_grid<GRID_NUM_WARPS, true, scalar_t, CELL_SIZE> :
+    //    device_a2a_adjacency_build_grid<GRID_NUM_WARPS, false, scalar_t, CELL_SIZE>;
+
+    //buildGridFn KERNEL_ARG2(gridSize, blockSize) (
+    //    ptrCounts,
     //    collapseResult.StridedMergeQuads.packed_accessor64<scalar_t, 3>(),
     //    gridCellsTensor.packed_accessor64<int32_t, 4>(),
     //    quadCellExtentsTensor.packed_accessor64<int32_t, 3>()
     //);
 
-    //device_a2a_adjacency_with_grid<GRID_NUM_WARPS, scalar_t> KERNEL_ARG3(gridSize, blockSize, smemSize) (
-    //    counts.data_ptr<int32_t>(),
+    //auto adjGridFn = counts.size(0) == 1 ?
+    //    device_a2a_adjacency_with_grid<GRID_NUM_WARPS, true, scalar_t> :
+    //    device_a2a_adjacency_with_grid<GRID_NUM_WARPS, false, scalar_t>;
+
+    //adjGridFn KERNEL_ARG3(gridSize, blockSize, smemSize) (
+    //    ptrCounts,
     //    iouThreshold,
     //    collapseResult.StridedMergeQuads.packed_accessor64<scalar_t, 3>(),
     //    gridCellsTensor.packed_accessor64<int32_t, 4>(),
@@ -1316,9 +1351,11 @@ AdjacencyResult compute_all_to_all_adjacency(
     gridSize = dim3{div_up(totalWork, blockSize.x),
                     static_cast<uint32_t>(counts.size(0))};
 
+    auto adjFn = counts.size(0) == 1 ? device_a2a_adjacency_sparse<true, scalar_t> : device_a2a_adjacency_sparse<false, scalar_t>;
+
     // This algorithm is O(n^2) with n being the current number of quads
-    device_a2a_adjacency_sparse<scalar_t> KERNEL_ARG2(gridSize, blockSize) (
-        counts.data_ptr<int32_t>(),
+    adjFn KERNEL_ARG2(gridSize, blockSize) (
+        ptrCounts,
         iouThreshold,
         collapseResult.StridedMergeQuads.packed_accessor64<scalar_t, 3>(),
         isStartTensor.packed_accessor64<bool, 2>(),
@@ -1328,9 +1365,7 @@ AdjacencyResult compute_all_to_all_adjacency(
 
 
 #ifdef NMS_VERIFY_CORRECTNESS
-    auto cpuCounts = counts.cpu();
-
-    cpu_a2a_adjacency_sparse<scalar_t>(cpuCounts.data_ptr<int32_t>(), iouThreshold,
+    cpu_a2a_adjacency_sparse<true>(ptrCounts, iouThreshold,
         collapseResult.StridedMergeQuads, cpuIsStartTensor, cpuAdjCountsTensor, cpuAdjValuesTensor);
 
     adjValuesTensor = std::get<0>(torch::sort(adjValuesTensor, /*dim=*/ 2));
@@ -1345,12 +1380,16 @@ AdjacencyResult compute_all_to_all_adjacency(
     auto maxDepthTensor = torch::tensor(0, adjCountsTensor.options());
 #endif
 
+    auto traverseFn = counts.size(0) == 1 ?
+                        device_flatten_graph_iterative<true> :
+                        device_flatten_graph_iterative<false>;
+
     blockSize = dim3{ 128, 1, 1 };
     gridSize = dim3{ 1, static_cast<uint32_t>(maxExCount), static_cast<uint32_t>(counts.size(0)) };
     smemSize = div_up(maxExCount * sizeof(visit_mask_t), sizeof(uint32_t)) * sizeof(uint32_t);
 
-    device_flatten_graph_iterative KERNEL_ARG3(gridSize, blockSize, smemSize) (
-        counts.data_ptr<int32_t>(),
+    traverseFn KERNEL_ARG3(gridSize, blockSize, smemSize) (
+        ptrCounts,
         isStartTensor.packed_accessor64<bool, 2>(),
         reinterpret_cast<uint32_t*>(adjCountsTensor.data_ptr<int32_t>()),
         reinterpret_cast<uint32_t*>(adjValuesTensor.data_ptr<int32_t>())
@@ -1360,7 +1399,7 @@ AdjacencyResult compute_all_to_all_adjacency(
     );
 
 #ifdef NMS_VERIFY_CORRECTNESS
-    cpu_flatten_graph(cpuCounts.data_ptr<int32_t>(), cpuIsStartTensor, cpuAdjCountsTensor, cpuAdjValuesTensor);
+    cpu_flatten_graph<true>(ptrCounts, cpuIsStartTensor, cpuAdjCountsTensor, cpuAdjValuesTensor);
 
     cpuAdjValuesTensor = std::get<0>(torch::sort(cpuAdjValuesTensor, /*dim=*/ 2));
     adjValuesTensor = std::get<0>(torch::sort(adjValuesTensor, /*dim=*/ 2));
@@ -1398,6 +1437,7 @@ AdjacencyResult compute_all_to_all_adjacency(
     cpuIsStartTensor = isStartTensor.cpu();
     cpuAdjCountsTensor = adjCountsTensor.cpu();
     cpuAdjValuesTensor = adjValuesTensor.cpu();
+    auto cpuCounts = counts.cpu();
     auto cpuCollapseIds = collapseResult.QuadIds.cpu();
 
     static std::vector<std::unordered_set<int32_t>> s_knownGroups;
@@ -1549,11 +1589,22 @@ nms_result_t
     dim3 blockSize(NUM_WARPS * 32, 1, 1);
     dim3 gridSize(1, adjResult.MaxExCount, counts.size(0));
 
+    // If the batch is only a single example, instead of hitting global memory for the count, we can
+    // just encode the count into the pointer instead
+    uint64_t ptrCounts = reinterpret_cast<uint64_t>(counts.data_ptr<int32_t>());
+    if (counts.size(0) == 1) {
+        ptrCounts = adjResult.MaxExCount;
+    }
+
     torch::Tensor outQuads = torch::empty({ numOutQuads, 4, 2 }, embedQuads.options());
     torch::Tensor outConf = torch::empty({ numOutQuads }, embedQuads.options());
 
-    device_a2a_collapse<NUM_WARPS, scalar_t>  KERNEL_ARG2(gridSize, blockSize) (
-        counts.packed_accessor64<int32_t, 1>(),
+    auto collapseFn = counts.size(0) == 1 ?
+        device_a2a_collapse<NUM_WARPS, scalar_t, true> :
+        device_a2a_collapse<NUM_WARPS, scalar_t, false>;
+
+    collapseFn KERNEL_ARG2(gridSize, blockSize) (
+        ptrCounts,
         embedQuads.packed_accessor64<scalar_t, 3>(),
         isLeadRow.packed_accessor64<bool, 2>(),
         regionCounts.data_ptr<int64_t>(),

nemotron-ocr/pyproject.toml

@@ -5,7 +5,6 @@ description = "Nemoton OCR"
 authors = [{ name = "NVIDIA Nemotron" }]
 requires-python = ">=3.12,<3.13"
 dependencies = [
-    "huggingface_hub>=0.20.0",
     "pandas>=2.3.3",
     "pillow>=12.0.0",
     "scikit-learn>=1.7.2",

nemotron-ocr/src/nemotron_ocr/inference/pipeline.py

@@ -6,7 +6,6 @@ import io
 import json
 import os
 from pathlib import Path
-from typing import Optional
 
 import numpy as np
 import torch
@@ -21,7 +20,6 @@ from nemotron_ocr.inference.post_processing.data.text_region import TextBlock
 from nemotron_ocr.inference.post_processing.quad_rectify import QuadRectify
 from nemotron_ocr.inference.post_processing.research_ops import parse_relational_results, reorder_boxes
 from nemotron_ocr.inference.pre_processing import interpolate_and_pad, pad_to_square
-from huggingface_hub import hf_hub_download
 from nemotron_ocr_cpp import quad_non_maximal_suppression, region_counts_to_indices, rrect_to_quads
 from PIL import Image, ImageDraw, ImageFont
 from torch import amp
@@ -39,57 +37,25 @@ MERGE_LEVELS = {"word", "sentence", "paragraph"}
 DEFAULT_MERGE_LEVEL = "paragraph"
 
 
-# HuggingFace repository for downloading model weights
-HF_REPO_ID = "nvidia/nemotron-ocr-v1"
-CHECKPOINT_FILES = ["detector.pth", "recognizer.pth", "relational.pth", "charset.txt"]
-
-
 class NemotronOCR:
     """
     A high-level pipeline for performing OCR on images.
-
-    Model weights are automatically downloaded from Hugging Face Hub
-    (nvidia/nemotron-ocr-v1) if not found locally.
     """
 
-    def __init__(self, model_dir: Optional[str] = None):
-        # If model_dir is provided and contains all required files, use it directly
-        if model_dir is not None:
-            local_path = Path(model_dir)
-            if all((local_path / f).is_file() for f in CHECKPOINT_FILES):
-                self._model_dir = local_path
-            else:
-                self._model_dir = self._download_checkpoints()
-        else:
-            self._model_dir = self._download_checkpoints()
+    def __init__(self, model_dir="./checkpoints"):
+        self._model_dir = Path(model_dir)
 
         self._load_models()
         self._load_charset()
         self._initialize_processors()
 
-    @staticmethod
-    def _download_checkpoints() -> Path:
-        """Download model checkpoints from HuggingFace Hub (cached locally after first download)."""
-        downloaded_path = None
-        for filename in CHECKPOINT_FILES:
-            downloaded_path = hf_hub_download(
-                repo_id=HF_REPO_ID,
-                filename=f"checkpoints/{filename}",
-            )
-        # All checkpoint files are in the same directory
-        return Path(downloaded_path).parent
-
     def _load_models(self):
         """Loads all necessary models into memory."""
         self.detector = FOTSDetector(coordinate_mode="RBOX", backbone="regnet_y_8gf", verbose=False)
-        self.detector.load_state_dict(
-            torch.load(self._model_dir / "detector.pth", weights_only=True), strict=True
-        )
+        self.detector.load_state_dict(torch.load(self._model_dir / "detector.pth"), strict=True)
 
         self.recognizer = TransformerRecognizer(nic=self.detector.num_features[-1], num_tokens=858, max_width=32)
-        self.recognizer.load_state_dict(
-            torch.load(self._model_dir / "recognizer.pth", weights_only=True), strict=True
-        )
+        self.recognizer.load_state_dict(torch.load(self._model_dir / "recognizer.pth"), strict=True)
 
         self.relational = GlobalRelationalModel(
             num_input_channels=self.detector.num_features,
@@ -98,9 +64,7 @@ class NemotronOCR:
             k=16,
             num_layers=4,
         )
-        self.relational.load_state_dict(
-            torch.load(self._model_dir / "relational.pth", weights_only=True), strict=True
-        )
+        self.relational.load_state_dict(torch.load(self._model_dir / "relational.pth"), strict=True)
 
         for model in (self.detector, self.recognizer, self.relational):
             model = model.cuda()
@@ -217,17 +181,29 @@ class NemotronOCR:
             e2e_det_conf = torch.sigmoid(det_conf)
             e2e_det_coords = rrect_to_quads(det_rboxes.float(), DETECTOR_DOWNSAMPLE)
 
-            quads, confidence, region_counts = quad_non_maximal_suppression(
-                e2e_det_coords,
-                e2e_det_conf,
-                prob_threshold=NMS_PROB_THRESHOLD,
-                iou_threshold=NMS_IOU_THRESHOLD,
-                kernel_height=2,
-                kernel_width=3,
-                max_regions=NMS_MAX_REGIONS,
-                verbose=False,
-            )[:3]
-
+            # FIXME: quad_non_maximal_suppression fails with batch size > 1
+            all_quads = []
+            all_confidence = []
+            all_region_counts = []
+
+            for idx in range(e2e_det_coords.shape[0]):
+                quads, confidence, region_counts = quad_non_maximal_suppression(
+                    e2e_det_coords[idx].unsqueeze(0),
+                    e2e_det_conf[idx].unsqueeze(0),
+                    prob_threshold=NMS_PROB_THRESHOLD,
+                    iou_threshold=NMS_IOU_THRESHOLD,
+                    kernel_height=2,
+                    kernel_width=3,
+                    max_regions=NMS_MAX_REGIONS,
+                    verbose=False,
+                )[:3]
+                all_quads.append(quads)
+                all_confidence.append(confidence)
+                all_region_counts.append(region_counts)
+
+            quads = torch.cat(all_quads, dim=0)
+            confidence = torch.cat(all_confidence, dim=0)
+            region_counts = torch.cat(all_region_counts, dim=0)
 
         if quads.shape[0] == 0:
             rec_rectified_quads = torch.empty(0, 128, 8, 32, dtype=torch.float32, device=padded_image.device)