Download checkpoint from HF hub in OcrReorderPipeline

app.py

@@ -8,21 +8,17 @@ import gradio as gr
 from inference import OcrReorderPipeline
 from transformers import AutoProcessor, LayoutLMv3Model, AutoTokenizer
 
-# ── 1) Load model/tokenizer/processor ─────────────────────────────────────
+# 1) Load model/tokenizer/processor
 repo      = "Uddipan107/ocr-layoutlmv3-base-t5-small"
 model     = LayoutLMv3Model.from_pretrained(repo)
 tokenizer = AutoTokenizer.from_pretrained(repo, subfolder="preprocessor")
 processor = AutoProcessor.from_pretrained(repo, subfolder="preprocessor", apply_ocr=False)
+pipe      = OcrReorderPipeline(model, tokenizer, processor, device=0)
 
-# instantiate your custom pipeline
-pipe = OcrReorderPipeline(model, tokenizer, processor, device=0)
-
-# ── 2) Inference function ─────────────────────────────────────────────────
 def infer(image_path, json_file):
-    # Extract filename
     img_name = os.path.basename(image_path)
 
-    # Parse NDJSON (one JSON object per line)
+    # Parse NDJSON
     data = []
     with open(json_file.name, "r", encoding="utf-8") as f:
         for line in f:
@@ -31,7 +27,6 @@ def infer(image_path, json_file):
                 continue
             data.append(json.loads(line))
 
-    # Find the matching entry
     entry = next((e for e in data if e["img_name"] == img_name), None)
     if entry is None:
         return f"❌ No JSON entry found for image '{img_name}'"
@@ -39,17 +34,15 @@ def infer(image_path, json_file):
     words = entry["src_word_list"]
     boxes = entry["src_wordbox_list"]
 
-    # Load & encode the image as base64
+    # Read & encode image
     img = Image.open(image_path).convert("RGB")
-    buf = BytesIO()
-    img.save(buf, format="PNG")
+    buf = BytesIO(); img.save(buf, format="PNG")
     b64 = base64.b64encode(buf.getvalue()).decode()
 
-    # ▶️ Correct keyword call so preprocess gets all args
-    reordered = pipe(image=b64, words=words, boxes=boxes)[0]
+    # ⚠️ Pass as `inputs`
+    reordered = pipe(inputs=b64, words=words, boxes=boxes)[0]
     return reordered
 
-# ── 3) Build the Gradio interface ───────────────────────────────────────────
 demo = gr.Interface(
     fn=infer,
     inputs=[
@@ -61,4 +54,5 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
+    # set share=True if you want a public link
     demo.launch()