Update app.py

app.py

@@ -17,10 +17,16 @@ from typing import Optional
 from PIL import Image
 from io import BytesIO
 import google.generativeai as genai
-import tempfile# Configure logging
+import tempfile
+
+# Configure logging
 logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)# Initialize FastAPI app
-app = FastAPI(title="Gemini Image Editing API with Razorpay")# Enable CORS for frontend
+logger = logging.getLogger(__name__)
+
+# Initialize FastAPI app
+app = FastAPI(title="Gemini Image Editing API with Razorpay")
+
+# Enable CORS for frontend
 app.add_middleware(
     CORSMiddleware,
     allow_origins=[
@@ -34,27 +40,45 @@ app.add_middleware(
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
-)# ===== API CONFIGURATION =====
+)
+
+# ===== API CONFIGURATION =====
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
-GEMINI_MODEL = "gemini-2.5-flash-image-preview"# Validate Gemini API key presence
+GEMINI_MODEL = "gemini-2.5-flash-image-preview"
+
+# Validate Gemini API key presence
 if not GEMINI_API_KEY:
     logger.error("GEMINI_API_KEY environment variable is not set")
-    raise ValueError("GEMINI_API_KEY environment variable is required")# Configure Gemini API
-genai.configure(api_key=GEMINI_API_KEY)# ===== RAZORPAY CONFIGURATION =====
+    raise ValueError("GEMINI_API_KEY environment variable is required")
+
+# Configure Gemini API
+genai.configure(api_key=GEMINI_API_KEY)
+
+# ===== RAZORPAY CONFIGURATION =====
 RAZORPAY_KEY_ID = os.getenv("RAZORPAY_KEY_ID")
 RAZORPAY_KEY_SECRET = os.getenv("RAZORPAY_KEY_SECRET")
-razorpay_client = razorpay.Client(auth=(RAZORPAY_KEY_ID, RAZORPAY_KEY_SECRET)) if RAZORPAY_KEY_ID and RAZORPAY_KEY_SECRET else None# ===== SUPABASE CONFIGURATION =====
+razorpay_client = razorpay.Client(auth=(RAZORPAY_KEY_ID, RAZORPAY_KEY_SECRET)) if RAZORPAY_KEY_ID and RAZORPAY_KEY_SECRET else None
+
+# ===== SUPABASE CONFIGURATION =====
 SUPABASE_URL = os.getenv("SUPABASE_URL")
 SUPABASE_KEY = os.getenv("SUPABASE_KEY")
-supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY) if SUPABASE_URL and SUPABASE_KEY else None# Pydantic models for JSON input validation
+supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY) if SUPABASE_URL and SUPABASE_KEY else None
+
+# Pydantic models for JSON input validation
 class CreateOrderRequest(BaseModel):
-    amount: intclass VerifyPaymentRequest(BaseModel):
+    amount: int
+
+class VerifyPaymentRequest(BaseModel):
     razorpay_order_id: str
     razorpay_payment_id: str
     razorpay_signature: str
-    user_id: Optional[str] = Noneclass GenerateImageRequest(BaseModel):
+    user_id: Optional[str] = None
+
+class GenerateImageRequest(BaseModel):
     prompt: str
-    user_id: Optional[str] = None# ===== AUTHENTICATION =====
+    user_id: Optional[str] = None
+
+# ===== AUTHENTICATION =====
 def generate_jwt_token():
     """Generate JWT token for API authentication"""
     payload = {
@@ -62,14 +86,18 @@ def generate_jwt_token():
         "exp": int(time.time()) + 1800,  # 30 minutes expiration
         "nbf": int(time.time()) - 5      # Not before 5 seconds ago
     }
-    return jwt.encode(payload, GEMINI_API_KEY, algorithm="HS256")# ===== IMAGE PROCESSING =====
+    return jwt.encode(payload, GEMINI_API_KEY, algorithm="HS256")
+
+# ===== IMAGE PROCESSING =====
 def prepare_image_base64(image_content: bytes):
     """Convert image bytes to base64 without prefix"""
     try:
         return base64.b64encode(image_content).decode('utf-8')
     except Exception as e:
         logger.error(f"Image processing failed: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Image processing failed: {str(e)}")def validate_image(image_content: bytes):
+        raise HTTPException(status_code=500, detail=f"Image processing failed: {str(e)}")
+
+def validate_image(image_content: bytes):
     """Validate image meets API requirements"""
     try:
         size_mb = len(image_content) / (1024 * 1024)
@@ -80,31 +108,36 @@ def prepare_image_base64(image_content: bytes):
             raise HTTPException(status_code=400, detail="Unsupported image format. Use PNG, JPG, JPEG, or WEBP")
         return True, ""
     except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Image validation error: {str(e)}")def generate_gemini_image(images: List[Image.Image], prompt: str):
+        raise HTTPException(status_code=400, detail=f"Image validation error: {str(e)}")
+
+def generate_gemini_image(images: List[Image.Image], prompt: str):
     """Generate or edit image using Gemini API"""
     try:
         contents = images + [prompt]
-        response = genai.GenerativeModel(GEMINI_MODEL).generate_content(contents)    text_response = ""
-    image_path = None
-
-    for part in response.candidates[0].content.parts:
-        if part.text:
-            text_response += part.text + "\n"
-        elif hasattr(part, 'inline_data') and part.inline_data:
-            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-                temp_path = tmp.name
-                generated_image = Image.open(BytesIO(part.inline_data.data))
-                generated_image.save(temp_path)
-                image_path = temp_path
-                logger.info(f"Generated image saved to: {temp_path} with prompt: {prompt}")
+        response = genai.GenerativeModel(GEMINI_MODEL).generate_content(contents)
+        text_response = ""
+        image_path = None
 
-    if image_path:
-        return image_path, ""
-    else:
-        return None, text_response.strip()
-except Exception as e:
-    logger.error(f"Gemini API error: {str(e)}")
-    raise HTTPException(status_code=500, detail=f"Gemini API error: {str(e)}")# ===== RAZORPAY FUNCTIONS =====
+        for part in response.candidates[0].content.parts:
+            if part.text:
+                text_response += part.text + "\n"
+            elif hasattr(part, 'inline_data') and part.inline_data:
+                with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+                    temp_path = tmp.name
+                    generated_image = Image.open(BytesIO(part.inline_data.data))
+                    generated_image.save(temp_path)
+                    image_path = temp_path
+                    logger.info(f"Generated image saved to: {temp_path} with prompt: {prompt}")
+
+        if image_path:
+            return image_path, ""
+        else:
+            return None, text_response.strip()
+    except Exception as e:
+        logger.error(f"Gemini API error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Gemini API error: {str(e)}")
+
+# ===== RAZORPAY FUNCTIONS =====
 def create_razorpay_order(amount: int):
     """Create a Razorpay order"""
     try:
@@ -122,7 +155,9 @@ def create_razorpay_order(amount: int):
         return order
     except Exception as e:
         logger.error(f"Failed to create Razorpay order: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Failed to create order: {str(e)}")def verify_payment_signature(order_id: str, payment_id: str, signature: str):
+        raise HTTPException(status_code=500, detail=f"Failed to create order: {str(e)}")
+
+def verify_payment_signature(order_id: str, payment_id: str, signature: str):
     """Verify Razorpay payment signature"""
     try:
         if not razorpay_client:
@@ -140,44 +175,50 @@ def create_razorpay_order(amount: int):
         return False
     except Exception as e:
         logger.error(f"Error verifying payment signature: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Verification error: {str(e)}")# ===== MAIN PROCESSING =====
+        raise HTTPException(status_code=500, detail=f"Verification error: {str(e)}")
+
+# ===== MAIN PROCESSING =====
 async def generate_image(images: List[bytes], prompt: str, user_id: Optional[str] = None):
     """Handle complete image generation workflow"""
     # Validate images
     for img_content in images:
         if img_content:
-            validate_image(img_content)# Convert bytes to PIL Images
-pil_images = []
-for img_content in images:
-    try:
-        img = Image.open(BytesIO(img_content))
-        if img.mode == "RGBA":
-            img = img.convert("RGBA")
-        pil_images.append(img)
-    except Exception as e:
-        logger.error(f"Image conversion failed: {str(e)}")
-        raise HTTPException(status_code=400, detail=f"Image conversion failed: {str(e)}")
+            validate_image(img_content)
 
-if len(pil_images) < 1:
-    raise HTTPException(status_code=400, detail="At least one image required")
+    # Convert bytes to PIL Images
+    pil_images = []
+    for img_content in images:
+        try:
+            img = Image.open(BytesIO(img_content))
+            if img.mode == "RGBA":
+                img = img.convert("RGBA")
+            pil_images.append(img)
+        except Exception as e:
+            logger.error(f"Image conversion failed: {str(e)}")
+            raise HTTPException(status_code=400, detail=f"Image conversion failed: {str(e)}")
 
-# Check user premium status if Supabase is configured
-if user_id and supabase:
-    try:
-        user_data = supabase.table("users").select("is_premium").eq("user_id", user_id).execute()
-        if not user_data.data or not user_data.data[0].get("is_premium"):
-            raise HTTPException(status_code=403, detail="Premium subscription required")
-    except Exception as e:
-        logger.error(f"Supabase user check failed: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"User verification failed: {str(e)}")
+    if len(pil_images) < 1:
+        raise HTTPException(status_code=400, detail="At least one image required")
+
+    # Check user premium status if Supabase is configured
+    if user_id and supabase:
+        try:
+            user_data = supabase.table("users").select("is_premium").eq("user_id", user_id).execute()
+            if not user_data.data or not user_data.data[0].get("is_premium"):
+                raise HTTPException(status_code=403, detail="Premium subscription required")
+        except Exception as e:
+            logger.error(f"Supabase user check failed: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"User verification failed: {str(e)}")
 
-# Generate image using Gemini API
-image_path, text_response = generate_gemini_image(pil_images, prompt)
+    # Generate image using Gemini API
+    image_path, text_response = generate_gemini_image(pil_images, prompt)
 
-if image_path:
-    return image_path
-else:
-    raise HTTPException(status_code=500, detail=f"Image generation failed: {text_response or 'Unknown error'}")# ===== API ENDPOINTS =====
+    if image_path:
+        return image_path
+    else:
+        raise HTTPException(status_code=500, detail=f"Image generation failed: {text_response or 'Unknown error'}")
+
+# ===== API ENDPOINTS =====
 @app.post("/generate")
 async def generate_image_endpoint(
     prompt: str = Form(...),
@@ -189,60 +230,66 @@ async def generate_image_endpoint(
         if len(images) < 1:
             raise HTTPException(status_code=400, detail="At least one image required")
         if len(images) > 4:
-            raise HTTPException(status_code=400, detail="Maximum 4 images allowed")    image_contents = [await image.read() for image in images]
-    output_path = await generate_image(image_contents, prompt, user_id)
-
-    return FileResponse(
-        path=output_path,
-        media_type="image/png",
-        filename=f"gemini_output_{Path(output_path).stem}.png"
-    )
-except HTTPException as e:
-    raise
-except Exception as e:
-    logger.error(f"Error in /generate: {str(e)}")
-    raise HTTPException(status_code=500, detail=str(e))@app.post("/create-razorpay-order")
+            raise HTTPException(status_code=400, detail="Maximum 4 images allowed")
+        image_contents = [await image.read() for image in images]
+        output_path = await generate_image(image_contents, prompt, user_id)
+
+        return FileResponse(
+            path=output_path,
+            media_type="image/png",
+            filename=f"gemini_output_{Path(output_path).stem}.png"
+        )
+    except HTTPException as e:
+        raise
+    except Exception as e:
+        logger.error(f"Error in /generate: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/create-razorpay-order")
 async def create_order_endpoint(
     request: Request,
     amount: Optional[int] = Form(None),
     body: Optional[CreateOrderRequest] = None
 ):
     """Create a Razorpay order (supports form-data and JSON)"""
-    logger.info("Received create order request")try:
-    if not razorpay_client:
-        raise HTTPException(status_code=500, detail="Razorpay configuration missing")
-    
-    # Handle JSON body if provided
-    if body and body.amount:
-        amount = body.amount
-    elif not amount:
-        try:
-            json_body = await request.json()
-            amount = json_body.get('amount')
-        except:
-            pass
-            
-    if not amount or amount <= 0:
-        raise HTTPException(status_code=422, detail="Missing or invalid 'amount' parameter")
-
-    logger.info(f"Creating order with amount: {amount}")
-    order = create_razorpay_order(amount)
-    
-    response_data = {
-        "id": order["id"],
-        "amount": order["amount"],
-        "currency": "INR",
-        "key_id": RAZORPAY_KEY_ID
-    }
-    
-    logger.info(f"Order created successfully: {order['id']}")
-    return JSONResponse(content=response_data)
-    
-except HTTPException:
-    raise
-except Exception as e:
-    logger.error(f"Error creating order: {str(e)}")
-    raise HTTPException(status_code=500, detail=f"Failed to create order: {str(e)}")@app.post("/verify-razorpay-payment")
+    logger.info("Received create order request")
+    try:
+        if not razorpay_client:
+            raise HTTPException(status_code=500, detail="Razorpay configuration missing")
+        
+        # Handle JSON body if provided
+        if body and body.amount:
+            amount = body.amount
+        elif not amount:
+            try:
+                json_body = await request.json()
+                amount = json_body.get('amount')
+            except:
+                pass
+                
+        if not amount or amount <= 0:
+            raise HTTPException(status_code=422, detail="Missing or invalid 'amount' parameter")
+
+        logger.info(f"Creating order with amount: {amount}")
+        order = create_razorpay_order(amount)
+        
+        response_data = {
+            "id": order["id"],
+            "amount": order["amount"],
+            "currency": "INR",
+            "key_id": RAZORPAY_KEY_ID
+        }
+        
+        logger.info(f"Order created successfully: {order['id']}")
+        return JSONResponse(content=response_data)
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error creating order: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Failed to create order: {str(e)}")
+
+@app.post("/verify-razorpay-payment")
 async def verify_payment_endpoint(
     request: Request,
     razorpay_order_id: Optional[str] = Form(None),
@@ -252,58 +299,61 @@ async def verify_payment_endpoint(
     body: Optional[VerifyPaymentRequest] = None
 ):
     """Verify Razorpay payment signature (supports form-data and JSON)"""
-    logger.info("Received payment verification request")try:
-    # Handle JSON body if provided
-    if body:
-        razorpay_order_id = razorpay_order_id or body.razorpay_order_id
-        razorpay_payment_id = razorpay_payment_id or body.razorpay_payment_id
-        razorpay_signature = razorpay_signature or body.razorpay_signature
-        user_id = user_id or body.user_id
-    else:
-        try:
-            json_body = await request.json()
-            razorpay_order_id = razorpay_order_id or json_body.get('razorpay_order_id')
-            razorpay_payment_id = razorpay_payment_id or json_body.get('razorpay_payment_id')
-            razorpay_signature = razorpay_signature or json_body.get('razorpay_signature')
-            user_id = user_id or json_body.get('user_id')
-        except:
-            pass
-
-    # Validate required fields
-    if not all([razorpay_order_id, razorpay_payment_id, razorpay_signature]):
-        missing_fields = []
-        if not razorpay_order_id: missing_fields.append("razorpay_order_id")
-        if not razorpay_payment_id: missing_fields.append("razorpay_payment_id")
-        if not razorpay_signature: missing_fields.append("razorpay_signature")
-        
-        logger.error(f"Missing required fields: {missing_fields}")
-        raise HTTPException(
-            status_code=422, 
-            detail=f"Missing required fields: {', '.join(missing_fields)}"
-        )
-
-    logger.info(f"Verifying payment for order_id: {razorpay_order_id}")
-    is_valid = verify_payment_signature(razorpay_order_id, razorpay_payment_id, razorpay_signature)
-    
-    if is_valid:
-        if user_id and supabase:
-            logger.info(f"Updating Supabase for user_id: {user_id}")
+    logger.info("Received payment verification request")
+    try:
+        # Handle JSON body if provided
+        if body:
+            razorpay_order_id = razorpay_order_id or body.razorpay_order_id
+            razorpay_payment_id = razorpay_payment_id or body.razorpay_payment_id
+            razorpay_signature = razorpay_signature or body.razorpay_signature
+            user_id = user_id or body.user_id
+        else:
             try:
-                supabase.table("users").update({"is_premium": True}).eq("user_id", user_id).execute()
-                logger.info(f"Successfully updated premium status for user: {user_id}")
-            except Exception as e:
-                logger.error(f"Failed to update Supabase: {str(e)}")
-        
-        return JSONResponse(content={"success": True, "message": "Payment verified successfully"})
-    else:
-        logger.warning(f"Payment verification failed for order: {razorpay_order_id}")
-        return JSONResponse(content={"success": False, "message": "Payment verification failed"}, status_code=400)
+                json_body = await request.json()
+                razorpay_order_id = razorpay_order_id or json_body.get('razorpay_order_id')
+                razorpay_payment_id = razorpay_payment_id or json_body.get('razorpay_payment_id')
+                razorpay_signature = razorpay_signature or json_body.get('razorpay_signature')
+                user_id = user_id or json_body.get('user_id')
+            except:
+                pass
+
+        # Validate required fields
+        if not all([razorpay_order_id, razorpay_payment_id, razorpay_signature]):
+            missing_fields = []
+            if not razorpay_order_id: missing_fields.append("razorpay_order_id")
+            if not razorpay_payment_id: missing_fields.append("razorpay_payment_id")
+            if not razorpay_signature: missing_fields.append("razorpay_signature")
+            
+            logger.error(f"Missing required fields: {missing_fields}")
+            raise HTTPException(
+                status_code=422, 
+                detail=f"Missing required fields: {', '.join(missing_fields)}"
+            )
+
+        logger.info(f"Verifying payment for order_id: {razorpay_order_id}")
+        is_valid = verify_payment_signature(razorpay_order_id, razorpay_payment_id, razorpay_signature)
         
-except HTTPException:
-    raise
-except Exception as e:
-    logger.error(f"Error verifying payment: {str(e)}")
-    raise HTTPException(status_code=500, detail=f"Verification error: {str(e)}")@app.get("/")
+        if is_valid:
+            if user_id and supabase:
+                logger.info(f"Updating Supabase for user_id: {user_id}")
+                try:
+                    supabase.table("users").update({"is_premium": True}).eq("user_id", user_id).execute()
+                    logger.info(f"Successfully updated premium status for user: {user_id}")
+                except Exception as e:
+                    logger.error(f"Failed to update Supabase: {str(e)}")
+            
+            return JSONResponse(content={"success": True, "message": "Payment verified successfully"})
+        else:
+            logger.warning(f"Payment verification failed for order: {razorpay_order_id}")
+            return JSONResponse(content={"success": False, "message": "Payment verification failed"}, status_code=400)
+            
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error verifying payment: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Verification error: {str(e)}")
+
+@app.get("/")
 async def index():
     return {
         "status": "Gemini Image Editing API with Razorpay is running",
@@ -312,13 +362,17 @@ async def index():
             "create_order": "POST /create-razorpay-order",
             "verify_payment": "POST /verify-razorpay-payment"
         }
-    }@app.get("/health")
+    }
+
+@app.get("/health")
 async def health_check():
     return {
         "status": "healthy",
         "razorpay_configured": bool(RAZORPAY_KEY_ID and RAZORPAY_KEY_SECRET),
         "supabase_configured": bool(SUPABASE_URL and SUPABASE_KEY),
         "gemini_configured": bool(GEMINI_API_KEY)
-    }if __name__ == "__main__":
+    }
+
+if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)