Update app.py

app.py

@@ -8,6 +8,8 @@ from PIL import Image
 from io import BytesIO
 import tempfile
 from dotenv import load_dotenv
+import warnings
+import io
 
 # Load environment variables from .env file
 load_dotenv()
@@ -23,101 +25,115 @@ def swap_clothing(person_image, clothing_image):
     Returns:
         The generated image with the clothing swapped and any relevant messages
     """
-    # Check if both images are provided
-    if person_image is None or clothing_image is None:
-        return None, "Please upload both images."
+    # Capture warnings in a string buffer
+    warning_buffer = io.StringIO()
+    warnings.filterwarnings('always')  # Ensure all warnings are shown
     
-    # 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 both uploaded images to temporary files
+    # Initialize variables outside the try block
     temp_files = []
-    try:
-        for img, prefix in [(person_image, "person"), (clothing_image, "clothing")]:
-            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
-                img.save(temp_file.name)
-                temp_files.append(temp_file.name)
-        
-        # Upload both files to Gemini
-        files = [
-            client.files.upload(file=temp_files[0]),  # person image
-            client.files.upload(file=temp_files[1]),  # clothing image
-        ]
-        
-        # Create the prompt
-        prompt = '''
-                    Edit the person's clothing by swapping it with the clothing in the clothing image.
-                    Retain the same facial features, pose and background from the person image.
-                    The output image should be an image of the person wearing the clothing from the clothing image.
-                    The image pose and background should be the same as the person image but with the new clothing:
-                '''
-        
-        contents = [
-            types.Content(
-                role="user",
-                parts=[
-                    types.Part.from_text(text="This is the person image. Do not change the face or features of the person. Pay attention and retain the background, pose, facial features."),
-                    types.Part.from_uri(
-                        file_uri=files[0].uri,
-                        mime_type=files[0].mime_type,
+    uploaded_files = []
+    client = None
+    output_image = None
+    output_text = ""
+    
+    with warnings.catch_warnings(record=True) as warning_list:
+        try:
+            # Check if both images are provided
+            if person_image is None or clothing_image is None:
+                return None, "Please upload both images."
+            
+            # 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."
+            
+            # Create a fresh client instance for each request
+            client = genai.Client(api_key=api_key)
+            
+            # Save both uploaded images to temporary files
+            for img, prefix in [(person_image, "person"), (clothing_image, "clothing")]:
+                with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
+                    img.save(temp_file.name)
+                    temp_files.append(temp_file.name)
+            
+            # Upload both files to Gemini with fresh file uploads
+            uploaded_files = [
+                client.files.upload(file=temp_files[0]),  # person image
+                client.files.upload(file=temp_files[1]),  # clothing image
+            ]
+            
+            # Create the prompt
+            prompt = '''
+                Edit the person's clothing by swapping it with the clothing in the clothing image.
+                Retain the same facial features, pose and background from the person image.
+                The output image should be an image of the person wearing the clothing from the clothing image.
+                The image pose and background should be the same as the person image but with the new clothing:
+            '''
+            
+            contents = [
+                types.Content(
+                    role="user",
+                    parts=[
+                        types.Part.from_text(text="This is the person image. Do not change the face or features of the person. Pay attention and retain the background, pose, facial features."),
+                        types.Part.from_uri(
+                            file_uri=uploaded_files[0].uri,
+                            mime_type=uploaded_files[0].mime_type,
+                        ),
+                        types.Part.from_text(text="This is the clothing image. Swap the clothing onto the person image."),
+                        types.Part.from_uri(
+                            file_uri=uploaded_files[1].uri,
+                            mime_type=uploaded_files[1].mime_type,
+                        ),
+                        types.Part.from_text(text=prompt),
+                        types.Part.from_uri(
+                            file_uri=uploaded_files[0].uri,
+                            mime_type=uploaded_files[0].mime_type,
+                        ),
+                    ],
+                ),
+            ]
+            
+            generate_content_config = types.GenerateContentConfig(
+                temperature=0.099,
+                top_p=0.95,
+                top_k=40,
+                max_output_tokens=8192,
+                response_modalities=[
+                    "image",
+                    "text",
+                ],
+                safety_settings=[
+                    types.SafetySetting(
+                        category="HARM_CATEGORY_HARASSMENT",
+                        threshold=HarmBlockThreshold.BLOCK_NONE,
+                    ),
+                    types.SafetySetting(
+                        category="HARM_CATEGORY_HATE_SPEECH",
+                        threshold=HarmBlockThreshold.BLOCK_NONE,
                     ),
-                    types.Part.from_text(text="This is the clothing image. Swap the clothing onto the person image."),
-                    types.Part.from_uri(
-                        file_uri=files[1].uri,
-                        mime_type=files[1].mime_type,
+                    types.SafetySetting(
+                        category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                        threshold=HarmBlockThreshold.BLOCK_NONE,
                     ),
-                    types.Part.from_text(text=prompt),
-                    types.Part.from_uri(
-                        file_uri=files[0].uri,
-                        mime_type=files[0].mime_type,
+                    types.SafetySetting(
+                        category="HARM_CATEGORY_DANGEROUS_CONTENT",
+                        threshold=HarmBlockThreshold.BLOCK_NONE,
                     ),
                 ],
-            ),
-        ]
-        
-        generate_content_config = types.GenerateContentConfig(
-            temperature=0.12,
-            top_p=0.95,
-            top_k=40,
-            max_output_tokens=8192,
-            response_modalities=[
-                "image",
-                "text",
-            ],
-            safety_settings=[
-                types.SafetySetting(
-                    category="HARM_CATEGORY_HARASSMENT",
-                    threshold=HarmBlockThreshold.BLOCK_NONE,
-                ),
-                types.SafetySetting(
-                    category="HARM_CATEGORY_HATE_SPEECH",
-                    threshold=HarmBlockThreshold.BLOCK_NONE,
-                ),
-                types.SafetySetting(
-                    category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
-                    threshold=HarmBlockThreshold.BLOCK_NONE,
-                ),
-                types.SafetySetting(
-                    category="HARM_CATEGORY_DANGEROUS_CONTENT",
-                    threshold=HarmBlockThreshold.BLOCK_NONE,
-                ),
-            ],
-            response_mime_type="text/plain",
-        )
+                response_mime_type="text/plain",
+            )
 
-        try:
             response = client.models.generate_content(
                 model="models/gemini-2.0-flash-exp",
                 contents=contents,
                 config=generate_content_config,
             )
-
-            output_image = None
-            output_text = ""
+            
+            # Add any warnings to the output text
+            if warning_list:
+                output_text += "\nWarnings:\n"
+                for warning in warning_list:
+                    output_text += f"- {warning.message}\n"
             
             # Process the response
             if response and hasattr(response, 'candidates') and response.candidates:
@@ -144,28 +160,34 @@ def swap_clothing(person_image, clothing_image):
                                 output_text += f"Error processing image: {str(img_error)}\n"
             else:
                 output_text = "The model did not generate a valid response. Please try again with different images."
-
-        except Exception as api_error:
-            output_text = f"API Error: {str(api_error)}\n\nDetails: {type(api_error).__name__}"
-            return None, output_text
+        
+        except Exception as e:
+            error_details = f"Error: {str(e)}\n\nType: {type(e).__name__}"
+            if warning_list:
+                error_details += "\n\nWarnings:\n"
+                for warning in warning_list:
+                    error_details += f"- {warning.message}\n"
+            print(f"Exception occurred: {error_details}")
+            return None, error_details
         
         finally:
-            # Clean up temporary files
+            # Clean up all temporary files
             for temp_file in temp_files:
                 if os.path.exists(temp_file):
                     os.unlink(temp_file)
+            
+            # Clean up any uploaded files if possible
+            for uploaded_file in uploaded_files:
+                try:
+                    if hasattr(client.files, 'delete') and uploaded_file:
+                        client.files.delete(uploaded_file.uri)
+                except:
+                    pass  # Best effort cleanup
+            
+            # Clear the client
+            client = None
         
         return output_image, output_text
-    
-    except Exception as e:
-        # Clean up temporary files in case of error
-        for temp_file in temp_files:
-            if os.path.exists(temp_file):
-                os.unlink(temp_file)
-        
-        error_details = f"Error: {str(e)}\n\nType: {type(e).__name__}"
-        print(f"Exception occurred: {error_details}")
-        return None, error_details
 
 # Create the Gradio interface
 def create_interface():