Update app.py

app.py

@@ -1,6 +1,10 @@
-from transformers import BlipProcessor, BlipForConditionalGeneration
+import os
+import shutil
+import base64
+import time
+import concurrent.futures
 from PIL import Image
-import gradio as gr
+from transformers import BlipProcessor, BlipForConditionalGeneration
 import torch
 from datetime import datetime
 from reportlab.lib.pagesizes import letter
@@ -8,11 +12,8 @@ from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image as PD
 from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
 from reportlab.lib import colors
 from simple_salesforce import Salesforce
-import os
 from dotenv import load_dotenv
-import base64
-import io
-import concurrent.futures
+import gradio as gr
 
 # Load environment variables from .env file
 load_dotenv()
@@ -30,23 +31,22 @@ except Exception as e:
     print(f"Failed to connect to Salesforce: {str(e)}")
 
 # Load BLIP model and processor
+device = "cuda" if torch.cuda.is_available() else "cpu"
 processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
 model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
-model.eval()
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+model.eval().to(device)
 
 # Inference function to generate captions dynamically based on image content
 def generate_captions_from_image(image):
     if image.mode != "RGB":
         image = image.convert("RGB")
-    
-    # Resize image for faster processing (use smaller resolution to speed up inference)
-    image = image.resize((320, 320))  # Reduced size for faster processing
+
+    # Resize for faster processing
+    image = image.resize((224, 224))  # Adjust to smaller resolution for faster inference
     
     # Preprocess the image and generate a caption
     inputs = processor(image, return_tensors="pt").to(device, torch.float16)
-    output = model.generate(**inputs, max_new_tokens=50)
+    output = model.generate(**inputs, max_length=50)
     caption = processor.decode(output[0], skip_special_tokens=True)
     
     return caption
@@ -139,9 +139,6 @@ def upload_file_to_salesforce(file_path, filename, sf_connection, file_type):
         
         # Generate a valid Salesforce URL for the ContentDocument
         content_document_url = f"https://{sf_connection.sf_instance}/sfc/servlet.shepherd/version/download/{content_version_id}"
-    
-
-        # Ensure the link is valid
         return content_document_id, content_document_url, f"File {filename} uploaded successfully"
     except Exception as e:
         return None, None, f"Error uploading {filename} to Salesforce: {str(e)}"
@@ -180,11 +177,6 @@ def generate_dpr(files):
     # Save DPR text to PDF
     pdf_result, pdf_filepath = save_dpr_to_pdf(dpr_output, image_paths, captions, pdf_filename)
     
-    salesforce_result = ""
-    pdf_content_document_id = None
-    pdf_url = None
-    image_content_document_ids = []
-    
     if sf and pdf_filepath:
         try:
             # Create Daily_Progress_Reports__c record
@@ -193,13 +185,11 @@ def generate_dpr(files):
                 'Detected_Activities__c': report_description  # Store in Detected_Activities__c field
             })
             dpr_record_id = dpr_record['id']
-            salesforce_result += f"Created Daily_Progress_Reports__c record with ID: {dpr_record_id}\n"
             
             # Upload PDF to Salesforce
             pdf_content_document_id, pdf_url, pdf_upload_result = upload_file_to_salesforce(
                 pdf_filepath, pdf_filename, sf, "pdf"
             )
-            salesforce_result += pdf_upload_result + "\n"
             
             # Link PDF to DPR record
             if pdf_content_document_id:
@@ -214,7 +204,6 @@ def generate_dpr(files):
                 sf.Daily_Progress_Reports__c.update(dpr_record_id, {
                     'PDF_URL__c': pdf_url  # Storing the PDF URL correctly
                 })
-                salesforce_result += f"Updated PDF URL for record ID {dpr_record_id}\n"
             
             # Upload images to Salesforce and link them to DPR record
             for file in files:
@@ -235,31 +224,40 @@ def generate_dpr(files):
                     sf.Daily_Progress_Reports__c.update(dpr_record_id, {
                         'Site_Images__c': image_content_document_id  # Storing the ContentDocumentId directly
                     })
-                    
-                    salesforce_result += image_upload_result + "\n"
                         
         except Exception as e:
-            salesforce_result += f"Error interacting with Salesforce: {str(e)}\n"
-    else:
-        salesforce_result = "Salesforce connection not available or PDF generation failed.\n"
+            pass  # No output for Salesforce errors now
     
-    # Return DPR text, PDF file, and Salesforce upload status
-    return (
-        dpr_output + f"\n\n{pdf_result}\n\nSalesforce Upload Status:\n{salesforce_result}",
-        pdf_filepath
-    )
+    # Return the PDF file for Gradio download (using shutil to copy and return the file)
+    if pdf_filepath:
+        # Copy the PDF file to a temporary directory for Gradio to serve it
+        temp_pdf_path = "/tmp/" + os.path.basename(pdf_filepath)
+        shutil.copy(pdf_filepath, temp_pdf_path)
+        
+        # Only return the DPR output and the PDF file path, excluding Salesforce upload details
+        return (
+            dpr_output + f"\n\n{pdf_result}",  # Removed Salesforce upload status
+            temp_pdf_path  # Returning the file path for download
+        )
+    else:
+        return (
+            dpr_output + f"\n\n{pdf_result}",  # Removed Salesforce upload status
+            None
+        )
+
 # Gradio interface for uploading multiple files, displaying DPR, and downloading PDF
 iface = gr.Interface(
     fn=generate_dpr,
     inputs=gr.Files(type="filepath", label="Upload Site Photos"),
     outputs=[
         gr.Textbox(label="Daily Progress Report"),
-        gr.File(label="Download PDF")
+        gr.File(label="Download PDF", interactive=False)
     ],
     title="Daily Progress Report Generator",
     description="Upload up to 10 site photos. The AI model will generate a text-based Daily Progress Report (DPR), save it as a PDF, and upload the PDF and images to Salesforce under Daily_Progress_Reports__c in the Files related list. Download the PDF locally if needed.",
-    allow_flagging="never"
+    allow_flagging="never",
+    css="#gradio-share-link-button-0 { display: none !important; }"
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()