Update app.py

app.py

@@ -8,6 +8,7 @@ import os
 from pathlib import Path
 from typing import List
 import io
+from PIL import Image
 import razorpay
 from razorpay.errors import SignatureVerificationError
 from supabase import create_client, Client
@@ -76,30 +77,29 @@ def validate_image(image_content: bytes):
         size_mb = len(image_content) / (1024 * 1024)
         if size_mb > 10:
             raise HTTPException(status_code=400, detail="Image too large (max 10MB)")
+        img = Image.open(io.BytesIO(image_content))
+        if img.format != "PNG":
+            raise HTTPException(status_code=400, detail="Only PNG images are supported")
         return True, ""
     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 multi-image generation task with Gemini API"""
+    """Create image generation task with Gemini API (single image or prompt-only)"""
     headers = {
         "Content-Type": "application/json"
     }
     subject_image_list = []
-    for img_content in subject_images:
-        if img_content:
-            base64_img = prepare_image_base64(img_content)
-            if base64_img:
-                subject_image_list.append({
-                    "inline_data": {
-                        "mime_type": "image/png",
-                        "data": base64_img
-                    }
-                })
-    
-    if len(subject_image_list) < 2:
-        raise HTTPException(status_code=400, detail="At least 2 subject images required")
+    if subject_images:  # Use only the first image if provided
+        base64_img = prepare_image_base64(subject_images[0])
+        if base64_img:
+            subject_image_list.append({
+                "inline_data": {
+                    "mime_type": "image/png",
+                    "data": base64_img
+                }
+            })
 
     payload = {
         "contents": [
@@ -136,6 +136,13 @@ async def generate_image(subject_images: List[bytes], prompt: str):
     
     task_response = create_multi_image_task(subject_images, prompt)
     try:
+        # Check for error in response
+        if "error" in task_response:
+            error_message = task_response["error"].get("message", "Unknown error")
+            error_code = task_response["error"].get("code", 500)
+            logger.error(f"API returned error: {error_code} - {error_message}")
+            raise HTTPException(status_code=500, detail=f"API error: {error_code} - {error_message}")
+        
         # Check response structure
         if "candidates" not in task_response or not task_response["candidates"]:
             logger.error(f"Invalid response structure: {task_response}")
@@ -147,12 +154,15 @@ async def generate_image(subject_images: List[bytes], prompt: str):
             raise HTTPException(status_code=500, detail="Invalid API response: No content parts found")
         
         parts = candidate["content"]["parts"]
+        logger.info(f"Response parts: {parts}")
         # Find the part with inline_data
         image_base64 = None
         for part in parts:
             if "inline_data" in part and "data" in part["inline_data"]:
                 image_base64 = part["inline_data"]["data"]
                 break
+            elif "text" in part:
+                logger.info(f"Text part found: {part['text']}")
         
         if not image_base64:
             logger.error(f"No inline_data found in response parts: {parts}")
@@ -174,14 +184,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(...)
+    images: List[UploadFile] = File(default=[])
 ):
-    """Endpoint to generate an image from multiple input images and a prompt"""
+    """Endpoint to generate an image from an optional input image and a prompt"""
     try:
-        if len(images) < 2:
-            raise HTTPException(status_code=400, detail="At least 2 images are required")
-        if len(images) > 4:
-            raise HTTPException(status_code=400, detail="Maximum 4 images allowed")
+        if len(images) > 1:
+            raise HTTPException(status_code=400, detail="Only one image is supported")
         image_contents = [await image.read() for image in images]
         output_path = await generate_image(image_contents, prompt)
         return FileResponse(