nv-ingest-api 2025.9.22.dev20250922__py3-none-any.whl → 2025.9.25.dev20250925__py3-none-any.whl

nv_ingest_api/internal/extract/image/chart_extractor.py

@@ -17,7 +17,8 @@ from nv_ingest_api.internal.schemas.extract.extract_chart_schema import ChartExt
 from nv_ingest_api.internal.schemas.meta.ingest_job_schema import IngestTaskChartExtraction
 from nv_ingest_api.util.image_processing.table_and_chart import join_yolox_graphic_elements_and_ocr_output
 from nv_ingest_api.util.image_processing.table_and_chart import process_yolox_graphic_elements
-from nv_ingest_api.internal.primitives.nim.model_interface.ocr import OCRModelInterface
+from nv_ingest_api.internal.primitives.nim.model_interface.ocr import PaddleOCRModelInterface
+from nv_ingest_api.internal.primitives.nim.model_interface.ocr import NemoRetrieverOCRModelInterface
 from nv_ingest_api.internal.primitives.nim.model_interface.ocr import get_ocr_model_name
 from nv_ingest_api.internal.primitives.nim import NimClient
 from nv_ingest_api.internal.primitives.nim.model_interface.yolox import YoloxGraphicElementsModelInterface
@@ -89,19 +90,12 @@ def _run_chart_inference(
     future_ocr_kwargs = dict(
         data=data_ocr,
         stage_name="chart_extraction",
-        max_batch_size=1 if ocr_client.protocol == "grpc" else 2,
         trace_info=trace_info,
     )
     if ocr_model_name == "paddle":
         future_ocr_kwargs.update(
             model_name="paddle",
-        )
-    elif ocr_model_name == "scene_text":
-        future_ocr_kwargs.update(
-            model_name=ocr_model_name,
-            input_names=["input", "merge_levels"],
-            dtypes=["FP32", "BYTES"],
-            merge_level="paragraph",
+            max_batch_size=1 if ocr_client.protocol == "grpc" else 2,
         )
     elif ocr_model_name == "scene_text_ensemble":
         future_ocr_kwargs.update(
@@ -134,7 +128,10 @@ def _run_chart_inference(
 
 
 def _validate_chart_inference_results(
-    yolox_results: Any, ocr_results: Any, valid_arrays: List[Any], valid_images: List[str]
+    yolox_results: Any,
+    ocr_results: Any,
+    valid_arrays: List[Any],
+    valid_images: List[str],
 ) -> Tuple[List[Any], List[Any]]:
     """
     Ensure inference results are lists and have expected lengths.
@@ -216,17 +213,12 @@ def _update_chart_metadata(
     return _merge_chart_results(base64_images, valid_indices, yolox_results, ocr_results, results)
 
 
-def _create_clients(
+def _create_yolox_client(
     yolox_endpoints: Tuple[str, str],
     yolox_protocol: str,
-    ocr_endpoints: Tuple[str, str],
-    ocr_protocol: str,
     auth_token: str,
-) -> Tuple[NimClient, NimClient]:
+) -> NimClient:
     yolox_model_interface = YoloxGraphicElementsModelInterface()
-    ocr_model_interface = OCRModelInterface()
-
-    logger.debug(f"Inference protocols: yolox={yolox_protocol}, ocr={ocr_protocol}")
 
     yolox_client = create_inference_client(
         endpoints=yolox_endpoints,
@@ -235,14 +227,29 @@ def _create_clients(
         infer_protocol=yolox_protocol,
     )
 
+    return yolox_client
+
+
+def _create_ocr_client(
+    ocr_endpoints: Tuple[str, str],
+    ocr_protocol: str,
+    ocr_model_name: str,
+    auth_token: str,
+) -> NimClient:
+    ocr_model_interface = (
+        NemoRetrieverOCRModelInterface() if ocr_model_name == "scene_text_ensemble" else PaddleOCRModelInterface()
+    )
+
     ocr_client = create_inference_client(
         endpoints=ocr_endpoints,
         model_interface=ocr_model_interface,
         auth_token=auth_token,
         infer_protocol=ocr_protocol,
+        enable_dynamic_batching=(True if ocr_model_name == "scene_text_ensemble" else False),
+        dynamic_batch_memory_budget_mb=32,
     )
 
-    return yolox_client, ocr_client
+    return ocr_client
 
 
 def extract_chart_data_from_image_internal(
@@ -285,13 +292,6 @@ def extract_chart_data_from_image_internal(
         return df_extraction_ledger, execution_trace_log
 
     endpoint_config = extraction_config.endpoint_config
-    yolox_client, ocr_client = _create_clients(
-        endpoint_config.yolox_endpoints,
-        endpoint_config.yolox_infer_protocol,
-        endpoint_config.ocr_endpoints,
-        endpoint_config.ocr_infer_protocol,
-        endpoint_config.auth_token,
-    )
 
     # Get the grpc endpoint to determine the model if needed
     ocr_grpc_endpoint = endpoint_config.ocr_endpoints[0]
@@ -334,6 +334,19 @@ def extract_chart_data_from_image_internal(
             base64_images.append(meta["content"])  # guaranteed by meets_criteria
 
         # 3) Call our bulk _update_metadata to get all results.
+        yolox_client = _create_yolox_client(
+            endpoint_config.yolox_endpoints,
+            endpoint_config.yolox_infer_protocol,
+            endpoint_config.auth_token,
+        )
+
+        ocr_client = _create_ocr_client(
+            endpoint_config.ocr_endpoints,
+            endpoint_config.ocr_infer_protocol,
+            ocr_model_name,
+            endpoint_config.auth_token,
+        )
+
         bulk_results = _update_chart_metadata(
             base64_images=base64_images,
             yolox_client=yolox_client,
@@ -356,13 +369,3 @@ def extract_chart_data_from_image_internal(
         logger.error("Error occurred while extracting chart data.", exc_info=True)
 
         raise
-
-    finally:
-        try:
-            if ocr_client is not None:
-                ocr_client.close()
-            if yolox_client is not None:
-                yolox_client.close()
-
-        except Exception as close_err:
-            logger.error(f"Error closing clients: {close_err}", exc_info=True)

nv_ingest_api/internal/extract/image/image_helpers/common.py

@@ -242,10 +242,6 @@ def extract_page_elements_from_images(
         logger.exception(f"Unhandled error during table/chart extraction: {str(e)}")
         raise
 
-    finally:
-        if yolox_client:
-            yolox_client.close()
-
     logger.debug(f"Extracted {len(page_elements)} tables and charts from image.")
     return page_elements
 

nv_ingest_api/internal/extract/image/infographic_extractor.py

@@ -12,11 +12,10 @@ from typing import Tuple
 import pandas as pd
 
 from nv_ingest_api.internal.primitives.nim import NimClient
-from nv_ingest_api.internal.primitives.nim.model_interface.ocr import OCRModelInterface
+from nv_ingest_api.internal.primitives.nim.model_interface.ocr import PaddleOCRModelInterface
+from nv_ingest_api.internal.primitives.nim.model_interface.ocr import NemoRetrieverOCRModelInterface
 from nv_ingest_api.internal.primitives.nim.model_interface.ocr import get_ocr_model_name
-from nv_ingest_api.internal.schemas.extract.extract_infographic_schema import (
-    InfographicExtractorSchema,
-)
+from nv_ingest_api.internal.schemas.extract.extract_infographic_schema import InfographicExtractorSchema
 from nv_ingest_api.util.image_processing.transforms import base64_to_numpy
 from nv_ingest_api.util.nim import create_inference_client
 from nv_ingest_api.util.image_processing.table_and_chart import reorder_boxes
@@ -101,19 +100,12 @@ def _update_infographic_metadata(
 
     infer_kwargs = dict(
         stage_name="infographic_extraction",
-        max_batch_size=1 if ocr_client.protocol == "grpc" else 2,
         trace_info=trace_info,
     )
     if ocr_model_name == "paddle":
         infer_kwargs.update(
             model_name="paddle",
-        )
-    elif ocr_model_name == "scene_text":
-        infer_kwargs.update(
-            model_name=ocr_model_name,
-            input_names=["input", "merge_levels"],
-            dtypes=["FP32", "BYTES"],
-            merge_level="paragraph",
+            max_batch_size=1 if ocr_client.protocol == "grpc" else 2,
         )
     elif ocr_model_name == "scene_text_ensemble":
         infer_kwargs.update(
@@ -144,25 +136,32 @@ def _update_infographic_metadata(
             # Each ocr_res is expected to be a tuple (text_predictions, bounding_boxes, conf_scores).
             ocr_res = reorder_boxes(*ocr_res)
 
-        results[original_index] = (base64_images[original_index], ocr_res[0], ocr_res[1])
+        results[original_index] = (
+            base64_images[original_index],
+            ocr_res[0],
+            ocr_res[1],
+        )
 
     return results
 
 
-def _create_clients(
+def _create_ocr_client(
     ocr_endpoints: Tuple[str, str],
     ocr_protocol: str,
+    ocr_model_name: str,
     auth_token: str,
 ) -> NimClient:
-    ocr_model_interface = OCRModelInterface()
-
-    logger.debug(f"Inference protocols: ocr={ocr_protocol}")
+    ocr_model_interface = (
+        NemoRetrieverOCRModelInterface() if ocr_model_name == "scene_text_ensemble" else PaddleOCRModelInterface()
+    )
 
     ocr_client = create_inference_client(
         endpoints=ocr_endpoints,
         model_interface=ocr_model_interface,
         auth_token=auth_token,
         infer_protocol=ocr_protocol,
+        enable_dynamic_batching=(True if ocr_model_name == "scene_text_ensemble" else False),
+        dynamic_batch_memory_budget_mb=32,
     )
 
     return ocr_client
@@ -239,11 +238,6 @@ def extract_infographic_data_from_image_internal(
         return df_extraction_ledger, execution_trace_log
 
     endpoint_config = extraction_config.endpoint_config
-    ocr_client = _create_clients(
-        endpoint_config.ocr_endpoints,
-        endpoint_config.ocr_infer_protocol,
-        endpoint_config.auth_token,
-    )
 
     # Get the grpc endpoint to determine the model if needed
     ocr_grpc_endpoint = endpoint_config.ocr_endpoints[0]
@@ -262,6 +256,13 @@ def extract_infographic_data_from_image_internal(
         base64_images = [df_extraction_ledger.at[idx, "metadata"]["content"] for idx in valid_indices]
 
         # Call bulk update to extract infographic data.
+        ocr_client = _create_ocr_client(
+            endpoint_config.ocr_endpoints,
+            endpoint_config.ocr_infer_protocol,
+            ocr_model_name,
+            endpoint_config.auth_token,
+        )
+
         bulk_results = _update_infographic_metadata(
             base64_images=base64_images,
             ocr_client=ocr_client,
@@ -283,6 +284,3 @@ def extract_infographic_data_from_image_internal(
         err_msg = "Error occurred while extracting infographic data."
         logger.exception(err_msg)
         raise
-
-    finally:
-        ocr_client.close()

nv_ingest_api/internal/extract/image/table_extractor.py

@@ -15,12 +15,13 @@ import pandas as pd
 
 from nv_ingest_api.internal.schemas.meta.ingest_job_schema import IngestTaskTableExtraction
 from nv_ingest_api.internal.enums.common import TableFormatEnum
-from nv_ingest_api.internal.primitives.nim.model_interface.ocr import OCRModelInterface
+from nv_ingest_api.internal.primitives.nim.model_interface.ocr import PaddleOCRModelInterface
+from nv_ingest_api.internal.primitives.nim.model_interface.ocr import NemoRetrieverOCRModelInterface
 from nv_ingest_api.internal.primitives.nim.model_interface.ocr import get_ocr_model_name
+from nv_ingest_api.internal.primitives.nim import NimClient
 from nv_ingest_api.internal.schemas.extract.extract_table_schema import TableExtractorSchema
 from nv_ingest_api.util.image_processing.table_and_chart import join_yolox_table_structure_and_ocr_output
 from nv_ingest_api.util.image_processing.table_and_chart import convert_ocr_response_to_psuedo_markdown
-from nv_ingest_api.internal.primitives.nim import NimClient
 from nv_ingest_api.internal.primitives.nim.model_interface.yolox import YoloxTableStructureModelInterface
 from nv_ingest_api.util.image_processing.transforms import base64_to_numpy
 from nv_ingest_api.util.nim import create_inference_client
@@ -31,7 +32,9 @@ PADDLE_MIN_WIDTH = 32
 PADDLE_MIN_HEIGHT = 32
 
 
-def _filter_valid_images(base64_images: List[str]) -> Tuple[List[str], List[np.ndarray], List[int]]:
+def _filter_valid_images(
+    base64_images: List[str],
+) -> Tuple[List[str], List[np.ndarray], List[int]]:
     """
     Filter base64-encoded images by their dimensions.
 
@@ -89,19 +92,12 @@ def _run_inference(
     future_ocr_kwargs = dict(
         data=data_ocr,
         stage_name="table_extraction",
-        max_batch_size=1 if ocr_client.protocol == "grpc" else 2,
         trace_info=trace_info,
     )
     if ocr_model_name == "paddle":
         future_ocr_kwargs.update(
             model_name="paddle",
-        )
-    elif ocr_model_name == "scene_text":
-        future_ocr_kwargs.update(
-            model_name=ocr_model_name,
-            input_names=["input", "merge_levels"],
-            dtypes=["FP32", "BYTES"],
-            merge_level="word",
+            max_batch_size=1 if ocr_client.protocol == "grpc" else 2,
         )
     elif ocr_model_name == "scene_text_ensemble":
         future_ocr_kwargs.update(
@@ -216,22 +212,22 @@ def _update_table_metadata(
     # Combine results with the original order.
     for idx, (yolox_res, ocr_res) in enumerate(zip(yolox_results, ocr_results)):
         original_index = valid_indices[idx]
-        results[original_index] = (base64_images[original_index], yolox_res, ocr_res[0], ocr_res[1])
+        results[original_index] = (
+            base64_images[original_index],
+            yolox_res,
+            ocr_res[0],
+            ocr_res[1],
+        )
 
     return results
 
 
-def _create_clients(
+def _create_yolox_client(
     yolox_endpoints: Tuple[str, str],
     yolox_protocol: str,
-    ocr_endpoints: Tuple[str, str],
-    ocr_protocol: str,
     auth_token: str,
-) -> Tuple[NimClient, NimClient]:
+) -> NimClient:
     yolox_model_interface = YoloxTableStructureModelInterface()
-    ocr_model_interface = OCRModelInterface()
-
-    logger.debug(f"Inference protocols: yolox={yolox_protocol}, ocr={ocr_protocol}")
 
     yolox_client = create_inference_client(
         endpoints=yolox_endpoints,
@@ -240,14 +236,29 @@ def _create_clients(
         infer_protocol=yolox_protocol,
     )
 
+    return yolox_client
+
+
+def _create_ocr_client(
+    ocr_endpoints: Tuple[str, str],
+    ocr_protocol: str,
+    ocr_model_name: str,
+    auth_token: str,
+) -> NimClient:
+    ocr_model_interface = (
+        NemoRetrieverOCRModelInterface() if ocr_model_name == "scene_text_ensemble" else PaddleOCRModelInterface()
+    )
+
     ocr_client = create_inference_client(
         endpoints=ocr_endpoints,
         model_interface=ocr_model_interface,
         auth_token=auth_token,
         infer_protocol=ocr_protocol,
+        enable_dynamic_batching=(True if ocr_model_name == "scene_text_ensemble" else False),
+        dynamic_batch_memory_budget_mb=32,
     )
 
-    return yolox_client, ocr_client
+    return ocr_client
 
 
 def extract_table_data_from_image_internal(
@@ -287,13 +298,6 @@ def extract_table_data_from_image_internal(
         return df_extraction_ledger, execution_trace_log
 
     endpoint_config = extraction_config.endpoint_config
-    yolox_client, ocr_client = _create_clients(
-        endpoint_config.yolox_endpoints,
-        endpoint_config.yolox_infer_protocol,
-        endpoint_config.ocr_endpoints,
-        endpoint_config.ocr_infer_protocol,
-        endpoint_config.auth_token,
-    )
 
     # Get the grpc endpoint to determine the model if needed
     ocr_grpc_endpoint = endpoint_config.ocr_endpoints[0]
@@ -335,6 +339,18 @@ def extract_table_data_from_image_internal(
         )
         enable_yolox = True if table_content_format in (TableFormatEnum.MARKDOWN,) else False
 
+        yolox_client = _create_yolox_client(
+            endpoint_config.yolox_endpoints,
+            endpoint_config.yolox_infer_protocol,
+            endpoint_config.auth_token,
+        )
+        ocr_client = _create_ocr_client(
+            endpoint_config.ocr_endpoints,
+            endpoint_config.ocr_infer_protocol,
+            ocr_model_name,
+            endpoint_config.auth_token,
+        )
+
         bulk_results = _update_table_metadata(
             base64_images=base64_images,
             yolox_client=yolox_client,
@@ -369,6 +385,3 @@ def extract_table_data_from_image_internal(
     except Exception:
         logger.exception("Error occurred while extracting table data.", exc_info=True)
         raise
-    finally:
-        yolox_client.close()
-        ocr_client.close()

nv_ingest_api/internal/extract/pdf/engines/pdfium.py

@@ -357,13 +357,6 @@ def _extract_page_elements(
     except Exception as e:
         logger.exception(f"Error in page element extraction: {str(e)}")
         raise
-    finally:
-        # Ensure client is closed properly
-        if yolox_client:
-            try:
-                yolox_client.close()
-            except Exception as e:
-                logger.warning(f"Error closing YOLOX client: {str(e)}")
 
     return extracted_page_elements
 

nv_ingest_api/internal/primitives/nim/__init__.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .nim_client import NimClient
+from .nim_client import get_nim_client_manager
 from .nim_model_interface import ModelInterface
 
-__all__ = ["NimClient", "ModelInterface"]
+__all__ = ["NimClient", "ModelInterface", "get_nim_client_manager"]

nv_ingest_api/internal/primitives/nim/model_interface/helpers.py

@@ -88,6 +88,8 @@ def preprocess_image_for_ocr(
     target_height: Optional[int] = None,
     target_width: Optional[int] = None,
     pad_how: str = "bottom_right",
+    normalize: bool = False,
+    channel_first: bool = False,
 ) -> np.ndarray:
     """
     Preprocesses an input image to be suitable for use with NemoRetriever-OCR.
@@ -122,10 +124,12 @@ def preprocess_image_for_ocr(
         how=pad_how,
     )
 
-    padded = padded / 255.0
+    if normalize:
+        padded = padded / 255.0
 
-    # NemoRetriever-OCR NIM (GRPC) requires input to be (channel, height, width).
-    transposed = padded.transpose((2, 0, 1))
+    if channel_first:
+        # NemoRetriever-OCR NIM (GRPC) requires input to be (channel, height, width).
+        padded = padded.transpose((2, 0, 1))
 
     # Metadata can used for inverting transformations on the resulting bounding boxes.
     metadata = {
@@ -137,7 +141,7 @@ def preprocess_image_for_ocr(
         "pad_width": pad_width,
     }
 
-    return transposed, metadata
+    return padded, metadata
 
 
 def is_ready(http_endpoint: str, ready_endpoint: str) -> bool: