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

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)