Update app.py

app.py

@@ -13,7 +13,7 @@ import os
 import glob
 import subprocess
 import tempfile
-import traceback
+import shutil
 from pathlib import Path
 
 import gradio as gr
@@ -44,21 +44,21 @@ def pdf_to_images(pdf_path: str, out_dir: str) -> list[str]:
 # Step 2 — Images → raw text (TrOCR)
 # ──────────────────────────────────────────────────────────────────────────────
 
-_trocr_processor = None
-_trocr_model = None
+_trocr_cache = {}
 
 def load_trocr(model_name: str):
-    global _trocr_processor, _trocr_model
-    if _trocr_processor is None:
+    global _trocr_cache
+    if model_name not in _trocr_cache:
         from transformers import TrOCRProcessor, VisionEncoderDecoderModel
-        _trocr_processor = TrOCRProcessor.from_pretrained(model_name)
-        _trocr_model = VisionEncoderDecoderModel.from_pretrained(model_name)
-        _trocr_model.eval()
-    return _trocr_processor, _trocr_model
+        processor = TrOCRProcessor.from_pretrained(model_name)
+        model = VisionEncoderDecoderModel.from_pretrained(model_name)
+        model.eval()
+        _trocr_cache[model_name] = (processor, model)
+    return _trocr_cache[model_name]
 
 
 def images_to_text(image_paths: list[str], model_name: str,
-                   progress: gr.Progress) -> list[str]:
+                    progress: gr.Progress) -> list[str]:
     import torch
     from PIL import Image
 
@@ -107,6 +107,8 @@ def refine_text(raw_pages: list[str], progress: gr.Progress) -> list[str]:
         chunks = [page[j:j+400] for j in range(0, len(page), 400)]
         corrected_chunks = []
         for chunk in chunks:
+            if not chunk.strip():
+                continue
             out = pipe(
                 f"grammar: {chunk}",
                 max_length=512,
@@ -114,7 +116,7 @@ def refine_text(raw_pages: list[str], progress: gr.Progress) -> list[str]:
                 early_stopping=True,
             )
             corrected_chunks.append(out[0]["generated_text"])
-        refined.append(" ".join(corrected_chunks))
+        refined.append(" ".join(corrected_chunks) if corrected_chunks else page)
 
     return refined
 
@@ -144,11 +146,10 @@ def text_to_docx(pages: list[str], output_path: str):
             doc.add_page_break()
         for line in page_text.split("\n"):
             line = line.strip()
+            p = doc.add_paragraph()
             if line:
-                p = doc.add_paragraph(line)
-                p.style.font.size = Pt(12)
-            else:
-                doc.add_paragraph("")
+                run = p.add_run(line)
+                run.font.size = Pt(12)
 
     doc.save(output_path)
 
@@ -166,31 +167,32 @@ def run_pipeline(
     if pdf_file is None:
         raise gr.Error("Please upload a PDF file first.")
 
+    # Handle Gradio filepath string vs object compatibility
+    pdf_path = pdf_file if isinstance(pdf_file, str) else pdf_file.name
+
     with tempfile.TemporaryDirectory() as tmp:
         # Step 1
         progress(0, desc="Converting PDF to images…")
-        images = pdf_to_images(pdf_file.name, tmp)
-        page_count = len(images)
+        images = pdf_to_images(pdf_path, tmp)
 
         # Step 2
-        progress(0, desc="Loading TrOCR model…")
+        progress(0.1, desc="Loading TrOCR model…")
         raw_pages = images_to_text(images, trocr_model, progress)
         raw_text = "\n\n[PAGE BREAK]\n\n".join(raw_pages)
 
         # Step 3
         if use_grammar:
-            progress(0, desc="Loading grammar correction model…")
+            progress(0.6, desc="Loading grammar correction model…")
             final_pages = refine_text(raw_pages, progress)
         else:
             final_pages = raw_pages
 
         # Step 4
-        progress(0.95, desc="Generating Word document…")
+        progress(0.9, desc="Generating Word document…")
         out_docx = os.path.join(tmp, "output.docx")
         text_to_docx(final_pages, out_docx)
 
         # Copy out of the temp dir so Gradio can serve it
-        import shutil
         dest = tempfile.NamedTemporaryFile(
             suffix=".docx", delete=False, prefix="scriptflow_"
         )
@@ -231,7 +233,7 @@ with gr.Blocks(title="ScriptFlow — Handwriting to Word", theme=gr.themes.Soft(
                 label="TrOCR model",
             )
             use_grammar = gr.Checkbox(
-                label="Apply T5 grammar correction (adds ~1 min per page)",
+                label="Apply T5 grammar correction (adds processing time)",
                 value=False,
             )
             run_btn = gr.Button("Convert", variant="primary")
@@ -245,9 +247,8 @@ with gr.Blocks(title="ScriptFlow — Handwriting to Word", theme=gr.themes.Soft(
         """
         ---
         **Notes**
-        - First run downloads models from HuggingFace (~330–440 MB) — subsequent runs use the cache.
+        - First run downloads models from Hugging Face — subsequent runs use the cache.
         - Grammar correction uses [`vennify/t5-base-grammar-correction`](https://huggingface.co/vennify/t5-base-grammar-correction).
-        - For best results use clear scans at ≥ 200 DPI.
         - Running on CPU: expect ~30–60 s per page for OCR.
         """
     )
@@ -259,6 +260,4 @@ with gr.Blocks(title="ScriptFlow — Handwriting to Word", theme=gr.themes.Soft(
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
-
-    
+    demo.launch()