Download checkpoint from HF hub in OcrReorderPipeline

app.py

@@ -8,17 +8,18 @@ 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)
 
+# ── 2) Inference function ──────────────────────────────────────────
 def infer(image_path, json_file):
     img_name = os.path.basename(image_path)
 
-    # Parse NDJSON
+    # Parse NDJSON entries from uploaded file
     data = []
     with open(json_file.name, "r", encoding="utf-8") as f:
         for line in f:
@@ -27,6 +28,7 @@ def infer(image_path, json_file):
                 continue
             data.append(json.loads(line))
 
+    # Find matching entry for this image
     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}'"
@@ -34,15 +36,17 @@ def infer(image_path, json_file):
     words = entry["src_word_list"]
     boxes = entry["src_wordbox_list"]
 
-    # Read & encode image
+    # Read and encode image to base64
     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()
 
-    # ⚠️ Pass as `inputs`
+    # Call pipeline with `inputs` keyword plus extra args
     reordered = pipe(inputs=b64, words=words, boxes=boxes)[0]
     return reordered
 
+# ── 3) Gradio interface ─────────────────────────────────────────────
 demo = gr.Interface(
     fn=infer,
     inputs=[