Update app.py

app.py

@@ -8,20 +8,13 @@ import os
 import time
 import jwt
 from pathlib import Path
-from typing import List, Optional
-from pydantic import BaseModel
+from typing import List
+import io
 
 # 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")
 
@@ -34,6 +27,12 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# ===== 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"
+
 # ===== AUTHENTICATION =====
 def generate_jwt_token():
     """Generate JWT token for API authentication"""
@@ -51,19 +50,17 @@ def prepare_image_base64(image_content: bytes):
         return base64.b64encode(image_content).decode('utf-8')
     except Exception as e:
         logger.error(f"Image processing failed: {str(e)}")
-        return None
+        raise HTTPException(status_code=500, detail=f"Image processing failed: {str(e)}")
 
 def validate_image(image_content: bytes):
     """Validate image meets API requirements"""
     try:
-        # 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)
+            raise HTTPException(status_code=400, detail="Image too large (max 10MB)")
         return True, ""
     except Exception as e:
-        return False, 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):
@@ -73,7 +70,6 @@ def create_multi_image_task(subject_images: List[bytes], prompt: str):
         "Content-Type": "application/json"
     }
     
-    # Prepare subject images list
     subject_image_list = []
     for img_content in subject_images:
         if img_content:
@@ -82,7 +78,7 @@ def create_multi_image_task(subject_images: List[bytes], prompt: str):
                 subject_image_list.append({"subject_image": base64_img})
     
     if len(subject_image_list) < 2:
-        return None, "At least 2 subject images required"
+        raise HTTPException(status_code=400, detail="At least 2 subject images required")
     
     payload = {
         "model_name": "kling-v2",
@@ -95,12 +91,12 @@ def create_multi_image_task(subject_images: List[bytes], prompt: str):
     try:
         response = requests.post(CREATE_TASK_ENDPOINT, json=payload, headers=headers)
         response.raise_for_status()
-        return response.json(), None
+        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}")
-        return None, f"API Error: {str(e)}"
+        raise HTTPException(status_code=500, detail=f"API Error: {str(e)}")
 
 def check_task_status(task_id: str):
     """Check task completion status"""
@@ -110,37 +106,26 @@ def check_task_status(task_id: str):
     try:
         response = requests.get(status_url, headers=headers)
         response.raise_for_status()
-        return response.json(), None
+        return response.json()
     except requests.exceptions.RequestException as e:
-        return None, f"Status check failed: {str(e)}"
+        raise HTTPException(status_code=500, detail=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
+            validate_image(img_content)
     
+    task_response = create_multi_image_task(subject_images, prompt)
     if task_response.get("code") != 0:
-        return None, f"API error: {task_response.get('message', 'Unknown error')}"
+        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}")
     
-    # Poll for results (max 10 minutes)
     for _ in range(60):
-        task_data, error = check_task_status(task_id)
-        if error:
-            return None, error
-            
+        task_data = check_task_status(task_id)
         status = task_data["data"]["task_status"]
         
         if status == "succeed":
@@ -148,62 +133,52 @@ async def generate_image(subject_images: List[bytes], prompt: str):
             try:
                 response = requests.get(image_url)
                 response.raise_for_status()
-                output_dir = Path("/tmp")  # Hugging Face Spaces uses /tmp for ephemeral storage
+                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 str(output_path), None
+                return output_path
             except Exception as e:
-                return None, f"Failed to download result: {str(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")
-            return None, f"Task failed: {error_msg}"
+            raise HTTPException(status_code=500, detail=f"Task failed: {error_msg}")
             
         time.sleep(10)
     
-    return None, "Task timed out after 10 minutes"
+    raise HTTPException(status_code=500, detail="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"
-    )
+    try:
+        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"kling_output_{Path(output_path).stem}.png"
+        )
+    except Exception as e:
+        logger.error(f"Error in /generate: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
 
-# Health check endpoint
-@app.get("/health")
-async def health_check():
-    return {"status": "healthy"}
+@app.get("/")
+async def index():
+    return {"status": "Kling AI Multi-Image Generator API is running"}
 
 if __name__ == "__main__":
     import uvicorn
-    # Hugging Face Spaces expects port 7860
     uvicorn.run(app, host="0.0.0.0", port=7860)