app.py

import gradio as gr
import requests
import base64
from pathlib import Path
import jwt
import time

# Kling AI API configuration (keys hardcoded as requested)
ACCESS_KEY_ID = "AGBGmadNd9hakFYfahytyQQJtN8CJmDJ"
ACCESS_KEY_SECRET = "dp3pAe4PpdmnAHCAPgEd3PyLmBQrkMde"
API_URL = "https://api-singapore.klingai.com/v1/predictions"

def generate_jwt_token():
    """Generate JWT token for Kling AI API authentication."""
    payload = {
        "iat": int(time.time()),
        "exp": int(time.time()) + 1800,  # Token expires in 30 minutes
        "access_key": ACCESS_KEY_ID
    }
    token = jwt.encode(payload, ACCESS_KEY_SECRET, algorithm="HS256")
    return token

def generate_image(image, prompt=""):
    """
    Call Kling AI API to generate an image from an input image.
    
    Args:
        image: Uploaded image file (from Gradio)
        prompt (str): Optional text prompt to guide image transformation
    
    Returns:
        str: Path to the generated image or error message
    """
    if not image:
        return "Error: Please upload a valid image (PNG/JPEG, <10 MB, ≥512x512 pixels)."
    
    # Convert image to base64
    try:
        with open(image, "rb") as img_file:
            image_base64 = base64.b64encode(img_file.read()).decode("utf-8")
    except Exception as e:
        return f"Error: Failed to process image. Details: {str(e)}"
    
    headers = {
        "Authorization": f"Bearer {generate_jwt_token()}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "inputs": [{"name": "input_image", "inputType": "URL", "url": image_base64}],
        "arguments": [
            {"name": "prompt", "value": prompt or "Transform the image"},
            {"name": "strength", "value": "0.7"},
            {"name": "output_format", "value": "png"},
            {"name": "task_type", "value": "image-to-image"}
        ]
    }
    
    try:
        response = requests.post(API_URL, json=payload, headers=headers, timeout=30)
        response.raise_for_status()
        
        data = response.json()
        task_id = data.get("task_id") or data.get("taskId")
        if not task_id:
            return "Error: No task ID returned by the API. Check endpoint or keys."
        
        # Poll for task completion
        status_url = f"https://api-singapore.klingai.com/v1/predictions/{task_id}"
        for _ in range(60):  # Poll for up to 5 minutes
            status_response = requests.get(status_url, headers=headers, timeout=30)
            status_response.raise_for_status()
            status_data = status_response.json()
            if status_data.get("status") == "succeeded":
                image_url = status_data.get("image_url") or status_data.get("result", {}).get("image_url")
                if not image_url:
                    return "Error: No image URL in API response."
                # Download the image
                image_response = requests.get(image_url, timeout=30)
                image_response.raise_for_status()
                output_path = Path("output_image.png")
                with open(output_path, "wb") as f:
                    f.write(image_response.content)
                return str(output_path)
            elif status_data.get("status") == "failed":
                return "Error: Image generation failed. Check image content (no NSFW, celebrities, or sensitive topics)."
            time.sleep(5)
        
        return "Error: Image generation timed out. Try a paid account if using free credits."
    
    except requests.exceptions.HTTPError as e:
        status_code = e.response.status_code
        if status_code == 401:
            return "Error: Authentication failed. Verify API keys."
        elif status_code == 404:
            return "Error: API endpoint not found. Check https://klingai.com/global/dev for correct endpoint."
        elif status_code == 422:
            return "Error: Invalid input. Use PNG/JPEG (<10 MB, ≥512x512 pixels) and a simple prompt."
        return f"Error: API request failed. Details: {str(e)}"
    except requests.exceptions.RequestException as e:
        return f"Error: Network issue. Ensure internet access and try again. Details: {str(e)}"

def chatbot_interface(image, prompt):
    """
    Gradio interface for image-to-image generation.
    
    Args:
        image: Uploaded image file
        prompt (str): Optional text prompt
    
    Returns:
        Image file path or error message
    """
    return generate_image(image, prompt)

# Define Gradio interface
iface = gr.Interface(
    fn=chatbot_interface,
    inputs=[
        gr.Image(type="filepath", label="Upload Image (PNG/JPEG, <10 MB, ≥512x512)"),
        gr.Textbox(lines=2, placeholder="Enter an optional prompt (e.g., 'Turn this into a cartoon')", label="Prompt")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Kling AI Image-to-Image Generator",
    description="Upload a PNG/JPEG image to generate a transformed image using Kling AI API. Ensure image is <10 MB, ≥512x512 pixels, and avoid NSFW or sensitive content."
)

# Launch the interface
if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)