Update app.py

app.py

@@ -51,66 +51,104 @@ if device.type == 'cuda':
 import traceback
 from concurrent.futures import ThreadPoolExecutor
 
-def handle_file(file, prompt, extra_prompt, max_new_tokens, progress=gr.Progress()):
+# def handle_file(file, prompt, extra_prompt, max_new_tokens, progress=gr.Progress()):
+#     try:
+#         file_path = file.name if hasattr(file, "name") else file
+#         filename = os.path.basename(file_path)
+#         ext = filename.lower().split('.')[-1]
+#         full_prompt = (prompt + "\n" + extra_prompt).strip() or ""
+
+#         print(f"[INFO] handle_file → {filename} (.{ext})")
+
+#         if ext == "pdf":
+#             try:
+#                 with open(file_path, "rb") as f:
+#                     pdf_bytes = f.read()
+#                 print(f"[INFO] Read PDF bytes: {len(pdf_bytes)} bytes")
+
+#                 doc = fitz.open(stream=pdf_bytes, filetype="pdf")
+#                 pages = []
+#                 zoom = DPI
+#                 mat = fitz.Matrix(zoom, zoom)
+#                 for i, page in enumerate(doc):
+#                     pix = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB)
+#                     img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+#                     if max(img.size) > 3072:
+#                         img.thumbnail((3072, 3072), Image.Resampling.LANCZOS)
+#                     pages.append(img)
+#                 print(f"[INFO] Converted PDF → {len(pages)} pages")
+
+#             except Exception as e:
+#                 traceback.print_exc()
+#                 return filename, f"[ERROR] PDF conversion failed: {e}"
+
+#             outputs = []
+#             with ThreadPoolExecutor(max_workers=4) as executor:
+#                 futures = [executor.submit(run_inference, img, full_prompt, max_new_tokens) for img in pages]
+#                 for idx, future in enumerate(futures):
+#                     try:
+#                         out = future.result()
+#                     except Exception as e:
+#                         traceback.print_exc()
+#                         out = f"[ERROR] Inference page {idx+1} failed: {e}"
+#                     outputs.append(out)
+#                     progress((idx) / len(pages), desc=f"Page {idx+1}/{len(pages)}")
+
+#             result = "\n\n--- Page Break ---\n\n".join(outputs)
+#             print("[INFO] handle_file done")
+#             return filename, result
+
+#         else:
+#             try:
+#                 img = Image.open(file_path)
+#                 print(f"[INFO] Opened image: {img.mode}, {img.size}")
+#             except Exception as e:
+#                 traceback.print_exc()
+#                 return filename, f"[ERROR] Image open failed: {e}"
+
+#             return filename, run_inference(img, full_prompt, max_new_tokens)
+
+#     except Exception as e:
+#         traceback.print_exc()
+#         return "error", f"[ERROR] handle_file unexpected: {e}"
+
+import time
+
+def handle_file(file, prompt, extra_prompt, max_new_tokens, progress=None):
     try:
         file_path = file.name if hasattr(file, "name") else file
         filename = os.path.basename(file_path)
         ext = filename.lower().split('.')[-1]
-        full_prompt = (prompt + "\n" + extra_prompt).strip() or ""
 
-        print(f"[INFO] handle_file → {filename} (.{ext})")
+        start = time.perf_counter()
 
         if ext == "pdf":
-            try:
-                with open(file_path, "rb") as f:
-                    pdf_bytes = f.read()
-                print(f"[INFO] Read PDF bytes: {len(pdf_bytes)} bytes")
-
-                doc = fitz.open(stream=pdf_bytes, filetype="pdf")
-                pages = []
-                zoom = DPI
-                mat = fitz.Matrix(zoom, zoom)
-                for i, page in enumerate(doc):
-                    pix = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB)
-                    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-                    if max(img.size) > 3072:
-                        img.thumbnail((3072, 3072), Image.Resampling.LANCZOS)
-                    pages.append(img)
-                print(f"[INFO] Converted PDF → {len(pages)} pages")
-
-            except Exception as e:
-                traceback.print_exc()
-                return filename, f"[ERROR] PDF conversion failed: {e}"
-
-            outputs = []
-            with ThreadPoolExecutor(max_workers=4) as executor:
-                futures = [executor.submit(run_inference, img, full_prompt, max_new_tokens) for img in pages]
-                for idx, future in enumerate(futures):
-                    try:
-                        out = future.result()
-                    except Exception as e:
-                        traceback.print_exc()
-                        out = f"[ERROR] Inference page {idx+1} failed: {e}"
-                    outputs.append(out)
-                    progress((idx) / len(pages), desc=f"Page {idx+1}/{len(pages)}")
-
-            result = "\n\n--- Page Break ---\n\n".join(outputs)
-            print("[INFO] handle_file done")
-            return filename, result
+            with open(file_path, "rb") as f:
+                pdf_bytes = f.read()
+
+            doc = fitz.open(stream=pdf_bytes, filetype="pdf")
+            scale = DPI / 72
+            mat = fitz.Matrix(scale, scale)
+            pages = []
+
+            for page in doc:
+                pix = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB)
+                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+                if max(img.size) > 3072:
+                    img.thumbnail((3072, 3072), Image.Resampling.LANCZOS)
+                pages.append(img)
 
         else:
-            try:
-                img = Image.open(file_path)
-                print(f"[INFO] Opened image: {img.mode}, {img.size}")
-            except Exception as e:
-                traceback.print_exc()
-                return filename, f"[ERROR] Image open failed: {e}"
+            img = Image.open(file_path)
+            if img.mode != "RGB":
+                img = img.convert("RGB")
 
-            return filename, run_inference(img, full_prompt, max_new_tokens)
+        end = time.perf_counter()
+        elapsed = end - start
+        return filename, f"OKE (time: {elapsed:.3f}s)"
 
     except Exception as e:
-        traceback.print_exc()
-        return "error", f"[ERROR] handle_file unexpected: {e}"
+        return "error", f"[ERROR] handle_file test failed: {e}"
 
 # def run_inference(img: Image.Image, prompt: str = "", max_new_tokens: int = 512) -> str:
 #     if img.mode != "RGB":