Update app.py

app.py

@@ -3,6 +3,7 @@ from PIL import Image
 import google.generativeai as genai
 import io
 import os
+import base64
 
 # Configure Google Gemini API (replace with your actual API key)
 GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
@@ -14,7 +15,6 @@ genai.configure(api_key=GOOGLE_API_KEY)
 # Function to generate the modified image
 def generate_modified_image(uploaded_image, background_description):
     try:
-        # Update the model name here
         model = genai.GenerativeModel('gemini-1.5-flash')
 
         img = Image.open(uploaded_image)
@@ -24,7 +24,7 @@ def generate_modified_image(uploaded_image, background_description):
             "Here is the image:",
             img,  # Pass the PIL Image object directly
             f"Modify the background of this image to: '{background_description}'. Be creative and make the new background look realistic and integrated with the foreground.",
-            "Output only the modified image."
+            "Output only the modified image. Ensure the output is a valid image format (like PNG or JPEG) encoded as bytes."
         ]
 
         response = model.generate_content(prompt_parts, stream=False)
@@ -32,20 +32,25 @@ def generate_modified_image(uploaded_image, background_description):
 
         if response and hasattr(response, 'parts') and len(response.parts) > 0:
             for part in response.parts:
-                if hasattr(part.data, 'image/png'):
-                    image_bytes = part.data['image/png']
-                    return Image.open(io.BytesIO(image_bytes))
-                elif hasattr(part.data, 'image/jpeg'):
-                    image_bytes = part.data['image/jpeg']
-                    return Image.open(io.BytesIO(image_bytes))
-            st.error("Generated content did not contain an image.")
+                try:
+                    # Try to get the image data directly if it's a blob
+                    if hasattr(part, 'blob'):
+                        image_bytes = part.blob.data
+                        return Image.open(io.BytesIO(image_bytes))
+                    elif hasattr(part, 'text'):
+                        st.warning(f"Received text response instead of image: {part.text}")
+                    else:
+                        st.warning(f"Unexpected part type in response: {part}")
+                except Exception as part_err:
+                    st.error(f"Error processing response part: {part_err}")
+            st.error("No valid image data found in the response.")
             return None
         else:
-            st.error("Failed to generate the modified image.")
+            st.error("Failed to get a valid response from the model.")
             return None
 
     except Exception as e:
-        st.error(f"An error occurred: {e}")
+        st.error(f"An error occurred during generation: {e}")
         return None
 
 # Streamlit web app