batched pipeline

app.py

@@ -26,7 +26,6 @@ import numpy as np
 from huggingface_hub import hf_hub_download
 from PIL import Image, ImageDraw
 
-from nemotron_ocr.inference.pipeline import NemotronOCR
 from nemotron_ocr.inference.pipeline_v2 import NemotronOCRV2
 
 MODELS = {
@@ -35,9 +34,7 @@ MODELS = {
     "v1 (legacy, English-only)": "v1",
 }
 
-PIPELINE_CHOICES = ["v2 (batched)", "v1 (original)"]
-
-_pipelines: dict[str, object] = {}
+_pipelines: dict[str, NemotronOCRV2] = {}
 
 GROUP_COLORS = [
     (76, 175, 80),
@@ -86,23 +83,14 @@ def _ensure_v1_model_dir() -> str:
     return model_dir
 
 
-def _get_pipeline(lang_key: str, pipeline_type: str):
-    cache_key = f"{lang_key}::{pipeline_type}"
-    if cache_key not in _pipelines:
-        use_v1_pipeline = pipeline_type == "v1 (original)"
-
+def _get_pipeline(lang_key: str) -> NemotronOCRV2:
+    if lang_key not in _pipelines:
         if lang_key in ("v1", "legacy"):
             model_dir = _ensure_v1_model_dir()
-            if use_v1_pipeline:
-                _pipelines[cache_key] = NemotronOCR(model_dir=model_dir)
-            else:
-                _pipelines[cache_key] = NemotronOCRV2(model_dir=model_dir)
+            _pipelines[lang_key] = NemotronOCRV2(model_dir=model_dir)
         else:
-            if use_v1_pipeline:
-                _pipelines[cache_key] = NemotronOCR(lang=lang_key)
-            else:
-                _pipelines[cache_key] = NemotronOCRV2(lang=lang_key)
-    return _pipelines[cache_key]
+            _pipelines[lang_key] = NemotronOCRV2(lang=lang_key)
+    return _pipelines[lang_key]
 
 
 def draw_boxes(image: Image.Image, predictions: list[dict]) -> Image.Image:
@@ -227,12 +215,12 @@ def format_text(predictions: list[dict], merge_level: str) -> str:
 
 
 @spaces.GPU(duration=120)
-def run_ocr(image: Image.Image, model_name: str, merge_level: str, pipeline_type: str):
+def run_ocr(image: Image.Image, model_name: str, merge_level: str):
     if image is None:
         return None, "Please upload an image."
 
     lang_key = MODELS[model_name]
-    ocr = _get_pipeline(lang_key, pipeline_type)
+    ocr = _get_pipeline(lang_key)
     img_array = np.array(image.convert("RGB"))
 
     if merge_level == "layout":
@@ -294,11 +282,6 @@ with gr.Blocks(
                 value="layout",
                 label="Output Mode",
             )
-            pipeline_type = gr.Radio(
-                choices=PIPELINE_CHOICES,
-                value="v2 (batched)",
-                label="Pipeline",
-            )
             run_btn = gr.Button("Run OCR", variant="primary")
 
         with gr.Column(scale=1):
@@ -311,7 +294,7 @@ with gr.Blocks(
 
     run_btn.click(
         fn=run_ocr,
-        inputs=[input_image, model_choice, merge_level, pipeline_type],
+        inputs=[input_image, model_choice, merge_level],
         outputs=[output_image, output_text],
     )
 

nemotron_ocr-1.0.0-cp312-cp312-linux_x86_64.whl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0ce2c3c3a382fdf90a2c9c5147fb91f9d6ccc516b312ae659ee42c97f99579ce
-size 45945424
+oid sha256:dc42583de574879c5c127d4e76398f04254bc5e3db651c3df0ccc942a2b48fa2
+size 45944873