Update app.py

app.py

@@ -114,15 +114,21 @@ from concurrent.futures import ThreadPoolExecutor
 
 import time
 
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
 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]
 
-        start = time.perf_counter()
+        full_prompt = (prompt + "\n" + extra_prompt).strip() or ""
+
+        start_total = time.perf_counter()
 
         if ext == "pdf":
+            start_convert = time.perf_counter()
             with open(file_path, "rb") as f:
                 pdf_bytes = f.read()
 
@@ -138,17 +144,55 @@ def handle_file(file, prompt, extra_prompt, max_new_tokens, progress=None):
                     img.thumbnail((3072, 3072), Image.Resampling.LANCZOS)
                 pages.append(img)
 
+            end_convert = time.perf_counter()
+            print(f"[INFO] Converted PDF → {len(pages)} pages in {(end_convert - start_convert):.3f}s")
+
+            # --- Xử lý inference đa luồng ---
+            start_infer = time.perf_counter()
+            outputs = []
+
+            def infer_page(img, idx):
+                # Thay run_inference bằng hàm inference thật hoặc giả
+                out = run_inference(img, full_prompt, max_new_tokens)
+                print(f"[DEBUG] Page {idx+1} inference done")
+                if progress:
+                    progress((idx) / len(pages), desc=f"Page {idx+1}/{len(pages)}")
+                return out
+
+            with ThreadPoolExecutor(max_workers=4) as executor:
+                futures = {executor.submit(infer_page, img, idx): idx for idx, img in enumerate(pages)}
+                for future in as_completed(futures):
+                    try:
+                        outputs.append(future.result())
+                    except Exception as e:
+                        outputs.append(f"[ERROR] Inference page failed: {e}")
+
+            end_infer = time.perf_counter()
+            print(f"[INFO] Inference all pages done in {(end_infer - start_infer):.3f}s")
+
+            total_time = end_infer - start_total
+            # Ghép kết quả các trang (thường là nối chuỗi, hoặc JSON array tuỳ model)
+            result = "\n\n--- Page Break ---\n\n".join(outputs)
+            return filename, f"OKE (total time: {total_time:.3f}s)\n{result}"
+
         else:
+            start_img = time.perf_counter()
             img = Image.open(file_path)
             if img.mode != "RGB":
                 img = img.convert("RGB")
+            end_img = time.perf_counter()
+            print(f"[INFO] Opened image in {(end_img - start_img):.3f}s")
 
-        end = time.perf_counter()
-        elapsed = end - start
-        return filename, f"OKE (time: {elapsed:.3f}s)"
+            start_infer = time.perf_counter()
+            result = run_inference(img, full_prompt, max_new_tokens)
+            end_infer = time.perf_counter()
+            total_time = end_infer - start_img
+            return filename, f"OKE (time: {total_time:.3f}s)\n{result}"
 
     except Exception as e:
-        return "error", f"[ERROR] handle_file test failed: {e}"
+        import traceback
+        traceback.print_exc()
+        return "error", f"[ERROR] handle_file failed: {e}"
 
 # def run_inference(img: Image.Image, prompt: str = "", max_new_tokens: int = 512) -> str:
 #     if img.mode != "RGB":