Create app.py

app.py

@@ -0,0 +1,196 @@
+import base64
+import os
+import gradio as gr
+from google import genai
+from google.genai import types
+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 swap_clothing(person_image, clothing_image):
+    """
+    Generate an image where the person from the first image is wearing clothing from the second image.
+    
+    Args:
+        person_image: The image containing the person
+        clothing_image: The image containing the clothing to swap
+    
+    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."
+    
+    # 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
+    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,
+                    ),
+                    types.Part.from_text(text="This is the clothing image."),
+                    types.Part.from_uri(
+                        file_uri=files[1].uri,
+                        mime_type=files[1].mime_type,
+                    ),
+                    types.Part.from_text(text=prompt),
+                    types.Part.from_uri(
+                        file_uri=files[0].uri,
+                        mime_type=files[0].mime_type,
+                    ),
+                ],
+            ),
+        ]
+        
+        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",
+        )
+
+        try:
+            response = client.models.generate_content(
+                model="models/gemini-2.0-flash-exp",
+                contents=contents,
+                config=generate_content_config,
+            )
+
+            output_image = None
+            output_text = ""
+            
+            # Process the response
+            if response and hasattr(response, 'candidates') and response.candidates:
+                candidate = response.candidates[0]
+                if hasattr(candidate, 'content') and candidate.content:
+                    for part in candidate.content.parts:
+                        if part.text is not None:
+                            output_text += part.text + "\n"
+                        elif part.inline_data is not None:
+                            try:
+                                if isinstance(part.inline_data.data, bytes):
+                                    image_data = part.inline_data.data
+                                else:
+                                    image_data = base64.b64decode(part.inline_data.data)
+                                
+                                with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+                                    temp_file.write(image_data)
+                                    temp_file_path = temp_file.name
+                                
+                                output_image = Image.open(temp_file_path)
+                                os.unlink(temp_file_path)
+                                
+                            except Exception as img_error:
+                                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
+        
+        finally:
+            # Clean up temporary files
+            for temp_file in temp_files:
+                if os.path.exists(temp_file):
+                    os.unlink(temp_file)
+        
+        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():
+    with gr.Blocks(title="Virtual Clothing Try-On") as app:
+        gr.Markdown("# Virtual Clothing Try-On")
+        gr.Markdown("Upload a photo of yourself and a photo of clothing you'd like to try on!")
+        
+        with gr.Row():
+            with gr.Column():
+                person_image = gr.Image(label="Your Photo", type="pil", image_mode="RGB")
+                clothing_image = gr.Image(label="Clothing Photo", type="pil", image_mode="RGB")
+                submit_btn = gr.Button("Generate")
+            
+            with gr.Column():
+                output_image = gr.Image(label="Result", type="pil")
+                output_text = gr.Textbox(label="Response", lines=3)
+        
+        submit_btn.click(
+            fn=swap_clothing,
+            inputs=[person_image, clothing_image],
+            outputs=[output_image, output_text]
+        )
+    
+    return app
+
+if __name__ == "__main__":
+    app = create_interface()
+    app.launch()