Update app.py

app.py

@@ -2,7 +2,7 @@ import logging
 from fastapi import FastAPI, UploadFile, File, HTTPException, Form, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse
-import requests
+import google.generativeai as genai
 import base64
 import os
 from pathlib import Path
@@ -21,7 +21,7 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Initialize FastAPI app
-app = FastAPI(title="Gemini 2.5 Flash Image Generator API with Razorpay")
+app = FastAPI(title="Gemini Image Generator API with Razorpay")
 
 # Enable CORS for the frontend
 app.add_middleware(
@@ -39,14 +39,18 @@ app.add_middleware(
 )
 
 # ===== API CONFIGURATION =====
-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}"
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+if not GEMINI_API_KEY:
+    logger.error("GEMINI_API_KEY is not set")
+    raise HTTPException(status_code=500, detail="GEMINI_API_KEY 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)
 
 # ===== 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))
+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")
@@ -64,14 +68,6 @@ class VerifyPaymentRequest(BaseModel):
     user_id: Optional[str] = None
 
 # ===== 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:
@@ -90,79 +86,39 @@ def validate_image(image_content: bytes):
 # ===== API FUNCTIONS =====
 def create_multi_image_task(subject_images: List[bytes], prompt: str):
     """Create image generation task with Gemini API (up to two images)"""
-    headers = {
-        "Content-Type": "application/json"
-    }
-    subject_image_list = []
-    for idx, img_content in enumerate(subject_images):
-        _, img_format = validate_image(img_content)
-        base64_img = prepare_image_base64(img_content)
-        if base64_img:
-            subject_image_list.append({
+    try:
+        model = genai.GenerativeModel(MODEL_NAME)
+        parts = []
+        for img_content in subject_images:
+            _, img_format = validate_image(img_content)
+            parts.append({
                 "inline_data": {
-                    "mime_type": f"image/{img_format}",
-                    "data": base64_img
+                    "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})
 
-    # Use a more descriptive prompt structure as per documentation
-    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."
-    
-    payload = {
-        "contents": [
-            {
-                "parts": [
-                    {"text": enhanced_prompt},
-                    *subject_image_list
-                ]
-            }
-        ],
-        "generationConfig": {
-            "response_mime_type": "image/png"
-        },
-        "safetySettings": [
-            {
-                "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"
-            }
-        ]
-    }
-
-    max_retries = 1
-    for attempt in range(max_retries + 1):
-        try:
-            logger.info(f"Sending request to Gemini API (attempt {attempt + 1}): {payload}")
-            response = requests.post(CREATE_TASK_ENDPOINT, json=payload, headers=headers)
-            response.raise_for_status()
-            data = response.json()
-            logger.info(f"API response: {data}")
-            if "safetyRatings" in data:
-                logger.info(f"Safety ratings: {data['safetyRatings']}")
-            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}")
-                if e.response.status_code in [429, 500] and attempt < max_retries:
-                    time.sleep(2 ** attempt)  # Exponential backoff
-                    continue
-            raise HTTPException(status_code=500, detail=f"API Error: {str(e)}")
+        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:
+        logger.error(f"API request failed: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"API Error: {str(e)}")
 
-# ===== MAIN PROCESSING =====
 async def generate_image(subject_images: List[bytes], prompt: str):
     """Handle complete image generation workflow"""
     if len(subject_images) != 2:
@@ -171,64 +127,23 @@ async def generate_image(subject_images: List[bytes], prompt: str):
     for img_content in subject_images:
         validate_image(img_content)
     
-    task_response = create_multi_image_task(subject_images, prompt)
+    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}")
-            raise HTTPException(status_code=500, detail="Invalid API response: No candidates found")
-        
-        candidate = task_response["candidates"][0]
-        if "content" not in candidate or "parts" not in candidate["content"]:
-            logger.error(f"Invalid content structure: {candidate}")
-            raise HTTPException(status_code=500, detail="Invalid API response: No content parts found")
-        
-        parts = candidate["content"]["parts"]
+        candidate = response.candidates[0]
+        parts = candidate.content.parts
         logger.info(f"Response parts: {parts}")
-        # Find the part with inline_data or file_uri
-        image_base64 = None
-        file_uri = None
-        text_response = None
+
+        image_data = None
         for part in parts:
-            if "inline_data" in part and "data" in part["inline_data"]:
-                image_base64 = part["inline_data"]["data"]
-                if not image_base64:
-                    logger.warning("Empty inline_data.data received")
-            elif "fileUri" in part:
-                file_uri = part["fileUri"]
-                logger.info(f"File URI found: {file_uri}")
-            elif "text" in part:
-                text_response = part["text"]
-                logger.info(f"Text part found: {text_response}")
-        
-        if not image_base64 and not file_uri:
-            error_detail = text_response or "No image data (inline_data or fileUri) found in API response"
-            if image_base64 == "":
-                error_detail = f"Empty inline_data.data returned by API: {text_response or 'No additional details'}"
-            logger.error(f"No image data in response parts: {parts}")
-            raise HTTPException(status_code=500, detail=f"API error: {error_detail}")
-        
-        if file_uri:
-            # Download image from file_uri
-            logger.info(f"Downloading image from {file_uri}")
-            response = requests.get(file_uri)
-            response.raise_for_status()
-            image_data = response.content
-        else:
-            # Decode base64 image
-            try:
-                image_data = base64.b64decode(image_base64)
-            except Exception as e:
-                logger.error(f"Failed to decode base64 image: {str(e)}")
-                raise HTTPException(status_code=500, detail=f"Failed to decode image data: {str(e)}")
-        
+            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")
+
         output_dir = Path("/tmp")
         output_dir.mkdir(exist_ok=True)
         output_path = output_dir / f"gemini_output_{int(time.time())}.png"
@@ -365,7 +280,7 @@ async def verify_payment_endpoint(
 @app.get("/")
 async def index():
     return {
-        "status": "Gemini 2.5 Flash Image Generator API with Razorpay is running",
+        "status": "Gemini Image Generator API with Razorpay is running",
         "endpoints": {
             "generate": "POST /generate",
             "create_order": "POST /create-razorpay-order",