Update app.py

app.py

@@ -85,17 +85,27 @@ def process_uploaded_file(file):
         
         if file is None:
             return "No file uploaded. Please upload a file."
+        
+        # Get the file path and extension
+        if isinstance(file, dict):
+            file_path = file["name"]
+        else:
+            file_path = file.name
             
-        file_path = file.name if isinstance(file, FileData) else file
-            
-        if file_path.lower().endswith('.pdf'):
+        # Get file extension
+        file_ext = file_path.lower().split('.')[-1]
+        
+        # Define allowed extensions
+        image_extensions = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'}
+        
+        if file_ext == 'pdf':
             doc_state.doc_type = 'pdf'
             try:
                 doc_state.current_doc_images, doc_state.current_doc_text = process_pdf_file(file_path)
                 return f"PDF processed successfully. Total pages: {len(doc_state.current_doc_images)}. You can now ask questions about the content."
             except Exception as e:
                 return f"Error processing PDF: {str(e)}. Please try a different PDF file."
-        else:
+        elif file_ext in image_extensions:
             doc_state.doc_type = 'image'
             try:
                 img = Image.open(file_path).convert("RGB")
@@ -108,6 +118,8 @@ def process_uploaded_file(file):
                 return "Image loaded successfully. You can now ask questions about the content."
             except Exception as e:
                 return f"Error processing image: {str(e)}. Please try a different image file."
+        else:
+            return f"Unsupported file type: {file_ext}. Please upload a PDF or image file (PNG, JPG, JPEG, GIF, BMP, WEBP)."
     except Exception as e:
         logger.error(f"Error in process_file: {str(e)}")
         return "An error occurred while processing the file. Please try again."
@@ -175,12 +187,12 @@ def clear_context():
 # Create the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Document Analyzer with Chat Support")
-    gr.Markdown("Upload a PDF or image and chat about its contents. For PDFs, all pages will be processed for visual analysis.")
+    gr.Markdown("Upload a PDF or image (PNG, JPG, JPEG, GIF, BMP, WEBP) and chat about its contents.")
     
     with gr.Row():
         file_upload = gr.File(
-            label="Upload Document (PDF or Image)",
-            file_types=["pdf", "image"]
+            label="Upload Document",
+            file_types=[".pdf", ".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp"]
         )
         upload_status = gr.Textbox(
             label="Upload Status",