Update app.py

app.py

@@ -7,6 +7,10 @@ from google.genai.types import HarmBlockThreshold
 from PIL import Image
 from io import BytesIO
 import tempfile
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
 
 def generate_image(input_image, partner_type):
     """
@@ -23,9 +27,12 @@ def generate_image(input_image, partner_type):
     if input_image is None:
         return None, "Please upload an image first."
     
-    client = genai.Client(
-        api_key=os.environ.get("GEMINI_API_KEY"),
-    )
+    # Get API key from environment variables
+    api_key = os.environ.get("GEMINI_API_KEY")
+    if not api_key:
+        return None, "GEMINI_API_KEY not found in environment variables."
+    
+    client = genai.Client(api_key=api_key)
     
     # Save the uploaded image to a temporary file
     with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
@@ -39,7 +46,7 @@ def generate_image(input_image, partner_type):
         ]
         
         # Create the prompt based on the partner type selection
-        prompt = f"add a {partner_type.lower()} to make them a couple. Do not change the environment or background."
+        prompt = f"add a {partner_type.lower()} beside the person in the picture. Do not change the environment or background."
         
         contents = [
             types.Content(
@@ -82,7 +89,7 @@ def generate_image(input_image, partner_type):
                 ),
                 types.SafetySetting(
                     category="HARM_CATEGORY_CIVIC_INTEGRITY",
-                    threshold=HarmBlockThreshold.BLOCK_NONE,
+                    threshold="OFF",
                 ),
             ],
             response_mime_type="text/plain",
@@ -106,9 +113,33 @@ def generate_image(input_image, partner_type):
                         if part.text is not None:
                             output_text += part.text + "\n"
                         elif part.inline_data is not None:
-                            # Decode the base64 data before passing to BytesIO
-                            image_data = base64.b64decode(part.inline_data.data)
-                            output_image = Image.open(BytesIO(image_data))
+                            try:
+                                # Add debug information
+                                # output_text += f"Received image data of type: {type(part.inline_data.data)}\n"
+                                # output_text += f"MIME type: {part.inline_data.mime_type}\n"
+                                
+                                # Try different approaches to handle the image data
+                                if isinstance(part.inline_data.data, bytes):
+                                    # If it's already bytes, use it directly
+                                    image_data = part.inline_data.data
+                                else:
+                                    # Otherwise, try to decode it from base64
+                                    image_data = base64.b64decode(part.inline_data.data)
+                                
+                                # Save to a temporary file first
+                                with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+                                    temp_file.write(image_data)
+                                    temp_file_path = temp_file.name
+                                
+                                # Open the image from the temporary file
+                                output_image = Image.open(temp_file_path)
+                                
+                                # Clean up the temporary file
+                                os.unlink(temp_file_path)
+                                
+                            except Exception as img_error:
+                                output_text += f"Error processing image: {str(img_error)}\n"
+                                output_text += f"Error type: {type(img_error).__name__}\n"
                 
                 # Check if there are finish details to display
                 if hasattr(candidate, 'finish_details') and candidate.finish_details: