Update app.py

app.py

@@ -16,14 +16,13 @@ def generate_modified_image(uploaded_image, background_description):
     try:
         model = genai.GenerativeModel('gemini-pro-vision')
 
-        image_parts = [
-            {"mime_type": uploaded_image.type, "data": uploaded_image.getvalue()}
-        ]
+        # Open the uploaded file as a PIL Image
+        img = Image.open(uploaded_image)
 
         prompt_parts = [
             "You are an AI that can modify the background of an image based on a text description.",
             "Here is the image:",
-            image_parts,
+            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."
         ]
@@ -32,8 +31,15 @@ def generate_modified_image(uploaded_image, background_description):
         response.resolve()
 
         if response and hasattr(response, 'parts') and len(response.parts) > 0:
-            image_bytes = response.parts[-1].data['image/png'] # Assuming the output is PNG
-            return Image.open(io.BytesIO(image_bytes))
+            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.")
+            return None
         else:
             st.error("Failed to generate the modified image.")
             return None