Update app.py

app.py

@@ -5,8 +5,6 @@ from fastapi.responses import FileResponse, JSONResponse
 import requests
 import base64
 import os
-import time
-import jwt
 from pathlib import Path
 from typing import List
 import io
@@ -21,17 +19,16 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Initialize FastAPI app
-app = FastAPI(title="Kling AI Multi-Image Generator API with Razorpay")
+app = FastAPI(title="Gemini 2.5 Flash Image Generator API with Razorpay")
 
-# Enable CORS for the frontend - Updated for Lovable domains
+# Enable CORS for the frontend
 app.add_middleware(
     CORSMiddleware,
     allow_origins=[
-        "https://hivili.web.app", 
+        "https://hivili.web.app",
         "http://localhost:3000",
         "https://*.lovable.dev",
         "https://*.sandbox.lovable.dev",
-        # Add pattern matching for Lovable subdomains
     ],
     allow_origin_regex=r"https://.*\.lovable\.dev|https://.*\.sandbox\.lovable\.dev",
     allow_credentials=True,
@@ -40,10 +37,9 @@ app.add_middleware(
 )
 
 # ===== API CONFIGURATION =====
-ACCESS_KEY_ID = os.getenv("ACCESS_KEY_ID", "AFyHfnQATghFdCMyAG3gRPbNY4TNKFGB")
-ACCESS_KEY_SECRET = os.getenv("ACCESS_KEY_SECRET", "TTepeLyBterLNM3brYPGmdndBnnyKJBA")
-API_BASE_URL = "https://api-singapore.klingai.com"
-CREATE_TASK_ENDPOINT = f"{API_BASE_URL}/v1/images/multi-image2image"
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "AIzaSyDL5Rilo7ptJpUOZdY6wy8PJYUcVcnDADs")
+API_BASE_URL = "https://generativelanguage.googleapis.com"
+CREATE_TASK_ENDPOINT = f"{API_BASE_URL}/v1beta/models/gemini-2.5-flash-image-preview:generateContent?key={GEMINI_API_KEY}"
 
 # ===== RAZORPAY CONFIGURATION =====
 RAZORPAY_KEY_ID = os.getenv("RAZORPAY_KEY_ID")
@@ -65,16 +61,6 @@ class VerifyPaymentRequest(BaseModel):
     razorpay_signature: str
     user_id: Optional[str] = None
 
-# ===== AUTHENTICATION =====
-def generate_jwt_token():
-    """Generate JWT token for API authentication"""
-    payload = {
-        "iss": ACCESS_KEY_ID,
-        "exp": int(time.time()) + 1800,  # 30 minutes expiration
-        "nbf": int(time.time()) - 5      # Not before 5 seconds ago
-    }
-    return jwt.encode(payload, ACCESS_KEY_SECRET, algorithm="HS256")
-
 # ===== IMAGE PROCESSING =====
 def prepare_image_base64(image_content: bytes):
     """Convert image bytes to base64 without prefix"""
@@ -96,88 +82,49 @@ def validate_image(image_content: bytes):
 
 # ===== API FUNCTIONS =====
 def create_multi_image_task(subject_images: List[bytes], prompt: str):
-    """Create multi-image generation task"""
+    """Create multi-image generation task with Gemini API"""
     headers = {
-        "Authorization": f"Bearer {generate_jwt_token()}",
         "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({"subject_image": 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")
-    
+
     payload = {
-        "model_name": "kling-v2",
-        "prompt": prompt,
-        "subject_image_list": subject_image_list,
-        "n": 1,
-        "aspect_ratio": "1:1"
+        "contents": [
+            {
+                "parts": [
+                    {"text": prompt},
+                    *subject_image_list
+                ]
+            }
+        ]
     }
-    
+
     try:
         response = requests.post(CREATE_TASK_ENDPOINT, json=payload, headers=headers)
         response.raise_for_status()
-        return response.json()
+        data = response.json()
+        if not data.get("candidates") or not data["candidates"][0].get("content"):
+            raise HTTPException(status_code=500, detail="No valid content returned from API")
+        return data
     except requests.exceptions.RequestException as e:
         logger.error(f"API request failed: {str(e)}")
         if hasattr(e, 'response') and e.response:
             logger.error(f"API response: {e.response.text}")
         raise HTTPException(status_code=500, detail=f"API Error: {str(e)}")
 
-def check_task_status(task_id: str):
-    """Check task completion status"""
-    headers = {"Authorization": f"Bearer {generate_jwt_token()}"}
-    status_url = f"{API_BASE_URL}/v1/images/multi-image2image/{task_id}"
-    
-    try:
-        response = requests.get(status_url, headers=headers)
-        response.raise_for_status()
-        return response.json()
-    except requests.exceptions.RequestException as e:
-        raise HTTPException(status_code=500, detail=f"Status check failed: {str(e)}")
-
-# ===== RAZORPAY FUNCTIONS =====
-def create_razorpay_order(amount: int):
-    """Create a Razorpay order"""
-    try:
-        if amount <= 0:
-            raise ValueError("Amount must be a positive integer")
-        order_data = {
-            "amount": amount * 100,  # Convert INR to paise
-            "currency": "INR",
-            "payment_capture": 1  # Auto-capture payment
-        }
-        order = razorpay_client.order.create(data=order_data)
-        logger.info(f"Razorpay order created successfully: {order['id']}")
-        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):
-    """Verify Razorpay payment signature"""
-    try:
-        params_dict = {
-            "razorpay_order_id": order_id,
-            "razorpay_payment_id": payment_id,
-            "razorpay_signature": signature
-        }
-        razorpay_client.utility.verify_payment_signature(params_dict)
-        logger.info(f"Payment signature verified successfully for order: {order_id}")
-        return True
-    except SignatureVerificationError as e:
-        logger.error(f"Payment signature verification failed: {str(e)}")
-        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 =====
 async def generate_image(subject_images: List[bytes], prompt: str):
     """Handle complete image generation workflow"""
@@ -186,37 +133,21 @@ async def generate_image(subject_images: List[bytes], prompt: str):
             validate_image(img_content)
     
     task_response = create_multi_image_task(subject_images, prompt)
-    if task_response.get("code") != 0:
-        raise HTTPException(status_code=500, detail=f"API error: {task_response.get('message', 'Unknown error')}")
-    
-    task_id = task_response["data"]["task_id"]
-    logger.info(f"Task created: {task_id}")
-    
-    for _ in range(60):
-        task_data = check_task_status(task_id)
-        status = task_data["data"]["task_status"]
+    try:
+        # Extract base64-encoded image from response
+        image_base64 = task_response["candidates"][0]["content"]["parts"][0]["inline_data"]["data"]
+        image_data = base64.b64decode(image_base64)
         
-        if status == "succeed":
-            image_url = task_data["data"]["task_result"]["images"][0]["url"]
-            try:
-                response = requests.get(image_url)
-                response.raise_for_status()
-                output_dir = Path("/tmp")
-                output_dir.mkdir(exist_ok=True)
-                output_path = output_dir / f"kling_output_{task_id}.png"
-                with open(output_path, "wb") as f:
-                    f.write(response.content)
-                return output_path
-            except Exception as e:
-                raise HTTPException(status_code=500, detail=f"Failed to download result: {str(e)}")
-                
-        elif status in ("failed", "canceled"):
-            error_msg = task_data["data"].get("task_status_msg", "Unknown error")
-            raise HTTPException(status_code=500, detail=f"Task failed: {error_msg}")
-            
-        time.sleep(10)
-    
-    raise HTTPException(status_code=500, detail="Task timed out after 10 minutes")
+        # Save the image to a temporary file
+        output_dir = Path("/tmp")
+        output_dir.mkdir(exist_ok=True)
+        output_path = output_dir / f"gemini_output_{int(time.time())}.png"
+        with open(output_path, "wb") as f:
+            f.write(image_data)
+        return output_path
+    except Exception as e:
+        logger.error(f"Failed to process API response: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Failed to process result: {str(e)}")
 
 # ===== API ENDPOINTS =====
 @app.post("/generate")
@@ -230,14 +161,12 @@ async def generate_image_endpoint(
             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")
-
         image_contents = [await image.read() for image in images]
         output_path = await generate_image(image_contents, prompt)
-
         return FileResponse(
             path=output_path,
             media_type="image/png",
-            filename=f"kling_output_{Path(output_path).stem}.png"
+            filename=f"gemini_output_{Path(output_path).stem}.png"
         )
     except Exception as e:
         logger.error(f"Error in /generate: {str(e)}")
@@ -251,39 +180,33 @@ async def create_order_endpoint(
 ):
     """Create a Razorpay order (supports form-data and JSON)"""
     logger.info("Received create order request")
-    
     if not RAZORPAY_KEY_ID or not RAZORPAY_KEY_SECRET:
         logger.error("Razorpay configuration missing")
         raise HTTPException(status_code=500, detail="Razorpay configuration missing")
     
-    # Handle JSON body if provided
     try:
         if body and body.amount:
             amount = body.amount
         elif not amount:
-            # Try to parse JSON from request body
             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": order["currency"],
             "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:
@@ -301,16 +224,13 @@ async def verify_payment_endpoint(
 ):
     """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 to parse JSON from request body
             try:
                 json_body = await request.json()
                 razorpay_order_id = razorpay_order_id or json_body.get('razorpay_order_id')
@@ -319,23 +239,23 @@ async def verify_payment_endpoint(
                 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")
-            
+            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, 
+                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}")
@@ -344,13 +264,10 @@ async def verify_payment_endpoint(
                     logger.info(f"Successfully updated premium status for user: {user_id}")
                 except Exception as e:
                     logger.error(f"Failed to update Supabase: {str(e)}")
-                    # Don't fail the payment verification if Supabase update fails
-            
             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:
@@ -360,7 +277,7 @@ async def verify_payment_endpoint(
 @app.get("/")
 async def index():
     return {
-        "status": "Kling AI Multi-Image Generator API with Razorpay is running",
+        "status": "Gemini 2.5 Flash Image Generator API with Razorpay is running",
         "endpoints": {
             "generate": "POST /generate",
             "create_order": "POST /create-razorpay-order",
@@ -376,6 +293,41 @@ async def health_check():
         "supabase_configured": bool(SUPABASE_URL and SUPABASE_KEY)
     }
 
+def create_razorpay_order(amount: int):
+    """Create a Razorpay order"""
+    try:
+        if amount <= 0:
+            raise ValueError("Amount must be a positive integer")
+        order_data = {
+            "amount": amount * 100,  # Convert INR to paise
+            "currency": "INR",
+            "payment_capture": 1  # Auto-capture payment
+        }
+        order = razorpay_client.order.create(data=order_data)
+        logger.info(f"Razorpay order created successfully: {order['id']}")
+        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):
+    """Verify Razorpay payment signature"""
+    try:
+        params_dict = {
+            "razorpay_order_id": order_id,
+            "razorpay_payment_id": payment_id,
+            "razorpay_signature": signature
+        }
+        razorpay_client.utility.verify_payment_signature(params_dict)
+        logger.info(f"Payment signature verified successfully for order: {order_id}")
+        return True
+    except SignatureVerificationError as e:
+        logger.error(f"Payment signature verification failed: {str(e)}")
+        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)}")
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)