Update app.py

app.py

@@ -14,6 +14,7 @@ from razorpay.errors import SignatureVerificationError
 from supabase import create_client, Client
 from pydantic import BaseModel
 from typing import Optional
+import time
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -75,28 +76,29 @@ def validate_image(image_content: bytes):
     """Validate image meets API requirements"""
     try:
         size_mb = len(image_content) / (1024 * 1024)
-        if size_mb > 10:
-            raise HTTPException(status_code=400, detail="Image too large (max 10MB)")
+        if size_mb > 20:
+            raise HTTPException(status_code=400, detail="Image too large (max 20MB)")
         img = Image.open(io.BytesIO(image_content))
-        if img.format != "PNG":
-            raise HTTPException(status_code=400, detail="Only PNG images are supported")
-        return True, ""
+        if img.format not in ["PNG", "JPEG"]:
+            raise HTTPException(status_code=400, detail="Only PNG or JPEG images are supported")
+        return True, img.format.lower()
     except Exception as e:
         raise HTTPException(status_code=400, detail=f"Image validation error: {str(e)}")
 
 # ===== API FUNCTIONS =====
 def create_multi_image_task(subject_images: List[bytes], prompt: str):
-    """Create image generation task with Gemini API (single image or prompt-only)"""
+    """Create image generation task with Gemini API (exactly two images)"""
     headers = {
         "Content-Type": "application/json"
     }
     subject_image_list = []
-    if subject_images:  # Use only the first image if provided
-        base64_img = prepare_image_base64(subject_images[0])
+    for idx, img_content in enumerate(subject_images):
+        _, img_format = validate_image(img_content)
+        base64_img = prepare_image_base64(img_content)
         if base64_img:
             subject_image_list.append({
                 "inline_data": {
-                    "mime_type": "image/png",
+                    "mime_type": f"image/{img_format}",
                     "data": base64_img
                 }
             })
@@ -130,9 +132,11 @@ def create_multi_image_task(subject_images: List[bytes], prompt: str):
 # ===== MAIN PROCESSING =====
 async def generate_image(subject_images: List[bytes], prompt: str):
     """Handle complete image generation workflow"""
+    if len(subject_images) != 2:
+        raise HTTPException(status_code=400, detail="Exactly two images are required")
+    
     for img_content in subject_images:
-        if img_content:
-            validate_image(img_content)
+        validate_image(img_content)
     
     task_response = create_multi_image_task(subject_images, prompt)
     try:
@@ -184,12 +188,12 @@ async def generate_image(subject_images: List[bytes], prompt: str):
 @app.post("/generate")
 async def generate_image_endpoint(
     prompt: str = Form(...),
-    images: List[UploadFile] = File(default=[])
+    images: List[UploadFile] = File(...)
 ):
-    """Endpoint to generate an image from an optional input image and a prompt"""
+    """Endpoint to generate an image from exactly two input images and a prompt"""
     try:
-        if len(images) > 1:
-            raise HTTPException(status_code=400, detail="Only one image is supported")
+        if len(images) != 2:
+            raise HTTPException(status_code=400, detail="Exactly two images are required")
         image_contents = [await image.read() for image in images]
         output_path = await generate_image(image_contents, prompt)
         return FileResponse(