Update app.py

app.py

@@ -2,35 +2,36 @@ import logging
 from fastapi import FastAPI, UploadFile, File, HTTPException, Form, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse
-import google.generativeai as genai
+import requests
 import base64
 import os
+import time
+import jwt
 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
 from pydantic import BaseModel
 from typing import Optional
-import time
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Initialize FastAPI app
-app = FastAPI(title="Gemini Image Generator API with Razorpay")
+app = FastAPI(title="Kling AI Multi-Image Generator API with Razorpay")
 
-# Enable CORS for the frontend
+# Enable CORS for the frontend - Updated for Lovable domains
 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,
@@ -39,18 +40,15 @@ app.add_middleware(
 )
 
 # ===== API CONFIGURATION =====
-GEMINI_API_KEY = os.getenv("AIzaSyDL5Rilo7ptJpUOZdY6wy8PJYUcVcnDADs")
-if not GEMINI_API_KEY:
-    logger.error("GEMINI_API_KEY is not set")
-    raise ValueError("GEMINI_API_KEY environment variable is not set")
-
-genai.configure(api_key=GEMINI_API_KEY)
-MODEL_NAME = "gemini-1.5-flash"  # Use a valid model (verify in Google's documentation)
+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"
 
 # ===== 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
+razorpay_client = razorpay.Client(auth=(RAZORPAY_KEY_ID, RAZORPAY_KEY_SECRET))
 
 # ===== SUPABASE CONFIGURATION =====
 SUPABASE_URL = os.getenv("SUPABASE_URL")
@@ -67,92 +65,158 @@ 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"""
+    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):
     """Validate image meets API requirements"""
     try:
         size_mb = len(image_content) / (1024 * 1024)
-        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 not in ["PNG", "JPEG"]:
-            raise HTTPException(status_code=400, detail="Only PNG or JPEG images are supported")
-        logger.info(f"Validated image: format={img.format}, size={size_mb:.2f}MB")
-        return True, img.format.lower()
+        if size_mb > 10:
+            raise HTTPException(status_code=400, detail="Image too large (max 10MB)")
+        return True, ""
     except Exception as e:
-        logger.error(f"Image validation error: {str(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 (up to two images)"""
+    """Create multi-image generation task"""
+    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})
+    
+    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"
+    }
+    
     try:
-        model = genai.GenerativeModel(MODEL_NAME)
-        parts = []
-        for img_content in subject_images:
-            _, img_format = validate_image(img_content)
-            parts.append({
-                "inline_data": {
-                    "data": base64.b64encode(img_content).decode('utf-8'),
-                    "mime_type": f"image/{img_format}"
-                }
-            })
-        enhanced_prompt = f"A photorealistic composition combining elements from the provided images: {prompt}. Ensure the scene is cohesive, with soft, natural lighting and a balanced aspect ratio of 16:9."
-        parts.append({"text": enhanced_prompt})
-
-        logger.info(f"Sending request to Gemini API with prompt: {prompt}")
-        response = model.generate_content(
-            parts,
-            generation_config={"response_mime_type": "image/png"},
-            safety_settings=[
-                {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
-                {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
-                {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
-                {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"}
-            ]
-        )
-        logger.info(f"API response: {response}")
-        if not response.candidates or not response.candidates[0].content:
-            raise HTTPException(status_code=500, detail="No valid content returned from API")
-        return response
-    except Exception as e:
+        response = requests.post(CREATE_TASK_ENDPOINT, json=payload, headers=headers)
+        response.raise_for_status()
+        return response.json()
+    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)}")
 
-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:
-        validate_image(img_content)
+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}"
     
-    response = create_multi_image_task(subject_images, prompt)
     try:
-        candidate = response.candidates[0]
-        parts = candidate.content.parts
-        logger.info(f"Response parts: {parts}")
+        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)}")
 
-        image_data = None
-        for part in parts:
-            if hasattr(part, 'inline_data') and part.inline_data.data:
-                image_data = part.inline_data.data
-                break
-            elif hasattr(part, 'text'):
-                logger.info(f"Text part found: {part.text}")
-
-        if not image_data:
-            raise HTTPException(status_code=500, detail="No image data found in API response")
+# ===== 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)}")
 
-        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
+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"Failed to process API response: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Failed to process result: {str(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"""
+    for img_content in subject_images:
+        if img_content:
+            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"]
+        
+        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")
 
 # ===== API ENDPOINTS =====
 @app.post("/generate")
@@ -160,16 +224,20 @@ async def generate_image_endpoint(
     prompt: str = Form(...),
     images: List[UploadFile] = File(...)
 ):
-    """Endpoint to generate an image from exactly two input images and a prompt"""
+    """Endpoint to generate an image from multiple input images and a prompt"""
     try:
-        if len(images) != 2:
-            raise HTTPException(status_code=400, detail="Exactly two images are required")
+        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")
+
         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"gemini_output_{Path(output_path).stem}.png"
+            filename=f"kling_output_{Path(output_path).stem}.png"
         )
     except Exception as e:
         logger.error(f"Error in /generate: {str(e)}")
@@ -183,33 +251,39 @@ 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:
@@ -227,13 +301,16 @@ 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')
@@ -242,23 +319,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}")
@@ -267,10 +344,13 @@ 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:
@@ -280,7 +360,7 @@ async def verify_payment_endpoint(
 @app.get("/")
 async def index():
     return {
-        "status": "Gemini Image Generator API with Razorpay is running",
+        "status": "Kling AI Multi-Image Generator API with Razorpay is running",
         "endpoints": {
             "generate": "POST /generate",
             "create_order": "POST /create-razorpay-order",
@@ -296,41 +376,6 @@ 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)