Update app.py

app.py

@@ -67,10 +67,9 @@ class FastPDFSearch:
                     img = img.convert('RGB')
                 img_byte_arr = io.BytesIO()
                 img.save(img_byte_arr, format='PNG')
-                img_byte_arr = img_byte_arr.getvalue()
-                img_base64 = base64.b64encode(img_byte_arr).decode()
+                img_byte_arr.seek(0)
                 doc.close()
-                return f"data:image/png;base64,{img_base64}"
+                return img_byte_arr
             doc.close()
             return None
         except Exception as e:
@@ -86,16 +85,19 @@ class FastPDFSearch:
         filename = result.metadata["filename"]
         page_number = result.metadata["page_number"]
         text = result.page_content
-        img = self.get_pdf_page_image(self.pdf_files[filename], page_number)
-        return text, img, f"{filename} - Page {page_number}"
+        img_bytes = self.get_pdf_page_image(self.pdf_files[filename], page_number)
+        if img_bytes:
+            return text, img_bytes, f"{filename} - Page {page_number}"
+        else:
+            return text, None, f"{filename} - Page {page_number}"
 
 # --- Gradio UI ---
 searcher = FastPDFSearch()
 
 def gradio_search(query):
-    text, img, label = searcher.search(query)
-    if img:
-        return text, [(img, label)]
+    text, img_bytes, label = searcher.search(query)
+    if img_bytes:
+        return text, [(img_bytes, label)]
     else:
         return text, []