Update app.py

app.py

@@ -4,7 +4,7 @@ import sys
 import tempfile
 
 from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import RedirectResponse, FileResponse
+from fastapi.responses import RedirectResponse, StreamingResponse
 import gradio as gr
 import requests
 import uvicorn
@@ -90,24 +90,37 @@ async def search(query: str, k: int):
 
     results = []
     for idx in top_k_indices:
-        results.append({"image": images[idx], "page": f"Page {idx}"})
+        img_byte_arr = BytesIO()
+        images[idx].save(img_byte_arr, format='PNG')
+        img_base64 = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
+        results.append({"image": img_base64, "page": f"Page {idx}"})
 
     # Generate PDF
     pdf_buffer = BytesIO()
     c = canvas.Canvas(pdf_buffer, pagesize=letter)
     width, height = letter
+
     for result in results:
-        img = result["image"]
+        img_base64 = result["image"]
+        img_data = base64.b64decode(img_base64)
+        
+        # Create a temporary file to hold the image
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
-            img.save(temp_file.name, format='PNG')
-            temp_file.seek(0)
+            temp_file.write(img_data)
+            temp_file.flush()
+            
+            # Draw the image from the temporary file
             c.drawImage(temp_file.name, 0, 0, width, height)
             c.showPage()
+        
+        # Clean up the temporary file
+        os.remove(temp_file.name)
 
     c.save()
     pdf_buffer.seek(0)
 
-    return FileResponse(pdf_buffer, media_type='application/pdf', filename='search_results.pdf')
+    # Use StreamingResponse to handle in-memory file
+    return StreamingResponse(pdf_buffer, media_type='application/pdf', filename='search_results.pdf')
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)