Update app.py

app.py

@@ -1,91 +1,209 @@
-# app.py
-from flask import Flask, request, jsonify
-import razorpay
+import logging
+from fastapi import FastAPI, UploadFile, File, HTTPException, Form
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse
+import requests
+import base64
 import os
-import hmac
-import hashlib
-import json
-
-app = Flask(__name__)
-
-# Get keys from Hugging Face Secrets (Settings -> Variables)
-RAZORPAY_KEY_ID = os.getenv("RAZORPAY_KEY_ID", "")
-RAZORPAY_KEY_SECRET = os.getenv("RAZORPAY_KEY_SECRET", "")
-
-if not RAZORPAY_KEY_ID or not RAZORPAY_KEY_SECRET:
-    print("WARNING: Razorpay keys not found in env. Set RAZORPAY_KEY_ID and RAZORPAY_KEY_SECRET in Space secrets.")
-
-client = razorpay.Client(auth=(RAZORPAY_KEY_ID, RAZORPAY_KEY_SECRET))
-
-def cors_json(resp):
-    # Helper to add CORS headers to Flask response dict
-    response = jsonify(resp)
-    response.headers.add("Access-Control-Allow-Origin", "*")
-    response.headers.add("Access-Control-Allow-Headers", "Content-Type")
-    return response
-
-@app.route("/create-order", methods=["POST", "OPTIONS"])
-def create_order():
-    # Preflight
-    if request.method == "OPTIONS":
-        response = app.response_class(status=204)
-        response.headers.add("Access-Control-Allow-Origin", "*")
-        response.headers.add("Access-Control-Allow-Headers", "Content-Type")
-        return response
-
-    data = request.get_json(force=True) or {}
-    amount = int(data.get("amount", 0))   # amount in rupees from frontend
-    if amount <= 0:
-        return cors_json({"error": "Invalid amount"}), 400
+import time
+import jwt
+from pathlib import Path
+from typing import List, Optional
+from pydantic import BaseModel
 
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ===== API CONFIGURATION =====
+# Load sensitive data from environment variables
+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"
+
+# Initialize FastAPI app
+app = FastAPI(title="Kling AI Multi-Image Generator API")
+
+# Enable CORS for the frontend
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["https://hivili.web.app"],  # Allow your frontend
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ===== 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:
-        options = {
-            "amount": amount * 100,  # paise
-            "currency": "INR",
-            "receipt": f"receipt_{os.urandom(4).hex()}",
-            "payment_capture": 1
-        }
-        order = client.order.create(options)
-        return cors_json(order)
+        return base64.b64encode(image_content).decode('utf-8')
     except Exception as e:
-        return cors_json({"error": str(e)}), 500
-
-@app.route("/verify-payment", methods=["POST", "OPTIONS"])
-def verify_payment():
-    # Verifies signature from Razorpay checkout
-    if request.method == "OPTIONS":
-        response = app.response_class(status=204)
-        response.headers.add("Access-Control-Allow-Origin", "*")
-        response.headers.add("Access-Control-Allow-Headers", "Content-Type")
-        return response
-
-    data = request.get_json(force=True) or {}
+        logger.error(f"Image processing failed: {str(e)}")
+        return None
+
+def validate_image(image_content: bytes):
+    """Validate image meets API requirements"""
     try:
-        razorpay_order_id = data.get("razorpay_order_id")
-        razorpay_payment_id = data.get("razorpay_payment_id")
-        razorpay_signature = data.get("razorpay_signature")
-
-        if not (razorpay_order_id and razorpay_payment_id and razorpay_signature):
-            return cors_json({"success": False, "message": "Missing fields"}), 400
-
-        payload = f"{razorpay_order_id}|{razorpay_payment_id}"
-        generated_signature = hmac.new(
-            RAZORPAY_KEY_SECRET.encode("utf-8"),
-            payload.encode("utf-8"),
-            hashlib.sha256
-        ).hexdigest()
-
-        if generated_signature == razorpay_signature:
-            return cors_json({"success": True})
-        else:
-            return cors_json({"success": False, "message": "Invalid signature"}), 400
+        # Check file size
+        size_mb = len(image_content) / (1024 * 1024)
+        if size_mb > 10:
+            return False, "Image too large (max 10MB)"
+        # Basic validation (add PIL-based dimension checks if needed)
+        return True, ""
     except Exception as e:
-        return cors_json({"success": False, "message": str(e)}), 500
+        return False, f"Image validation error: {str(e)}"
+
+# ===== API FUNCTIONS =====
+def create_multi_image_task(subject_images: List[bytes], prompt: str):
+    """Create multi-image generation task"""
+    headers = {
+        "Authorization": f"Bearer {generate_jwt_token()}",
+        "Content-Type": "application/json"
+    }
+    
+    # Prepare subject images list
+    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:
+        return None, "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:
+        response = requests.post(CREATE_TASK_ENDPOINT, json=payload, headers=headers)
+        response.raise_for_status()
+        return response.json(), None
+    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}")
+        return None, 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(), None
+    except requests.exceptions.RequestException as e:
+        return None, f"Status check failed: {str(e)}"
+
+# ===== MAIN PROCESSING =====
+async def generate_image(subject_images: List[bytes], prompt: str):
+    """Handle complete image generation workflow"""
+    # Validate images
+    for img_content in subject_images:
+        if img_content:
+            is_valid, error_msg = validate_image(img_content)
+            if not is_valid:
+                return None, error_msg
+    
+    # Create task
+    task_response, error = create_multi_image_task(subject_images, prompt)
+    if error:
+        return None, error
+    
+    if task_response.get("code") != 0:
+        return None, f"API error: {task_response.get('message', 'Unknown error')}"
+    
+    task_id = task_response["data"]["task_id"]
+    logger.info(f"Task created: {task_id}")
+    
+    # Poll for results (max 10 minutes)
+    for _ in range(60):
+        task_data, error = check_task_status(task_id)
+        if error:
+            return None, error
+            
+        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")  # Hugging Face Spaces uses /tmp for ephemeral storage
+                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 str(output_path), None
+            except Exception as e:
+                return None, f"Failed to download result: {str(e)}"
+                
+        elif status in ("failed", "canceled"):
+            error_msg = task_data["data"].get("task_status_msg", "Unknown error")
+            return None, f"Task failed: {error_msg}"
+            
+        time.sleep(10)
+    
+    return None, "Task timed out after 10 minutes"
+
+# ===== API ENDPOINTS =====
+class GenerateImageRequest(BaseModel):
+    prompt: str
+
+@app.post("/generate")
+async def generate_image_endpoint(
+    prompt: str = Form(...),
+    images: List[UploadFile] = File(...)
+):
+    """Endpoint to generate an image from multiple input images and a prompt"""
+    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")
+
+    # Read image contents
+    image_contents = []
+    for image in images:
+        content = await image.read()
+        image_contents.append(content)
+
+    # Generate image
+    output_path, error = await generate_image(image_contents, prompt)
+    if error:
+        raise HTTPException(status_code=500, detail=error)
+
+    # Return the generated image file
+    return FileResponse(
+        path=output_path,
+        media_type="image/png",
+        filename=f"kling_output_{Path(output_path).stem}.png"
+    )
 
-@app.route("/health", methods=["GET"])
-def health():
-    return cors_json({"ok": True, "app": "razorpay-backend"})
+# Health check endpoint
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}
 
 if __name__ == "__main__":
-    # Hugging Face Spaces runs on port 7860
-    app.run(host="0.0.0.0", port=7860)
+    import uvicorn
+    # Hugging Face Spaces expects port 7860
+    uvicorn.run(app, host="0.0.0.0", port=7860)