Update app.py

app.py

@@ -3,8 +3,11 @@ import fitz
 import torch
 from transformers import pipeline
 import time, logging, re
+import matplotlib
+matplotlib.use('Agg')  # Use non-interactive backend for headless environments
 import matplotlib.pyplot as plt
 import io
+from PIL import Image
 
 logging.basicConfig(level=logging.ERROR)
 device = -1  # CPU-only
@@ -20,7 +23,7 @@ def visualize_chunk_status(chunk_data):
     status_colors = {'summarized': 'green', 'skipped': 'orange', 'error': 'red'}
     labels = [f"C{i['chunk']}" for i in chunk_data]
     colors = [status_colors.get(i['status'], 'gray') for i in chunk_data]
-    times = [i.get('time', 0.1) for i in chunk_data]  # Avoid zero-time bars
+    times = [i.get('time', 0.1) for i in chunk_data]
 
     fig, ax = plt.subplots(figsize=(10, 2.5))
     ax.barh(labels, times, color=colors)
@@ -31,7 +34,8 @@ def visualize_chunk_status(chunk_data):
     buf = io.BytesIO()
     plt.savefig(buf, format='png')
     buf.seek(0)
-    return buf
+    plt.close(fig)  # Release memory
+    return Image.open(buf)
 
 def summarize_file(file_bytes):
     start = time.time()
@@ -56,7 +60,7 @@ def summarize_file(file_bytes):
 
     for i, chunk in enumerate(chunks):
         chunk_start = time.time()
-        chunk_result = {'chunk': i+1, 'status': '', 'time': 0}
+        chunk_result = {'chunk': i + 1, 'status': '', 'time': 0}
 
         if time.time() - start > 20:
             summaries.append("⚠️ Stopped early")
@@ -78,15 +82,15 @@ def summarize_file(file_bytes):
         chunk_info.append(chunk_result)
 
     final_summary = f"**Chars**: {len(text)}\n**Time**: {time.time()-start:.2f}s\n\n" + "\n\n".join(summaries)
-    image_buf = visualize_chunk_status(chunk_info)
-    return final_summary, image_buf
+    image = visualize_chunk_status(chunk_info)
+    return final_summary, image
 
 demo = gr.Interface(
     fn=summarize_file,
     inputs=gr.File(label="📄 Upload PDF", type="binary"),
     outputs=[
         gr.Textbox(label="📝 Summarized Output"),
-        gr.Image(label="📊 Visual Process Flow")
+        gr.Image(label="📊 Visual Process Flow", type="pil")
     ],
     title="AI-Powered PDF Summarizer",
     description="Summarizes long PDFs (up to 300,000 characters) and visualizes chunk-level automation status."
@@ -97,3 +101,4 @@ if __name__ == "__main__":
         demo.launch(share=False, server_port=7860)
     except Exception as e:
         print(f"❌ Gradio launch failed: {str(e)}")
+