app.py

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

# === Kling AI API configuration ===
ACCESS_KEY_ID = "AGBGmadNd9hakFYfahytyQQJtN8CJmDJ"
ACCESS_KEY_SECRET = "dp3pAe4PpdmnAHCAPgEd3PyLmBQrkMde"
API_URL = "https://api-singapore.klingai.com/v1/generate/image"

# === Generate JWT Token ===
def generate_jwt_token():
    headers = {
        "alg": "HS256",
        "typ": "JWT"
    }
    now = int(time.time())
    payload = {
        "iss": ACCESS_KEY_ID,
        "exp": now + 1800,  # valid for 30 mins
        "nbf": now - 5
    }
    return jwt.encode(payload, ACCESS_KEY_SECRET, algorithm="HS256", headers=headers)

# === Image to Image Generation ===
def generate_image(reference_image, prompt=""):
    if not reference_image:
        return None, None

    try:
        with open(reference_image, "rb") as img_file:
            reference_base64 = base64.b64encode(img_file.read()).decode("utf-8")
    except Exception as e:
        return None, f"Error reading reference image: {str(e)}"

    headers = {
        "Authorization": f"Bearer {generate_jwt_token()}",
        "Content-Type": "application/json"
    }

    payload = {
        "reference_image": reference_base64,
        "mode": "face",  # face reference mode
        "prompt": prompt or "Match face with high fidelity",
        "strength": 0.97,
        "output_format": "png"
    }

    try:
        # Start generation
        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("id")
        if not task_id:
            return None, f"Error: No task ID returned: {data}"

        # Poll for completion
        status_url = f"https://api-singapore.klingai.com/v1/predictions/{task_id}"
        for _ in range(60):  # up to 5 minutes
            status_response = requests.get(status_url, headers=headers, timeout=30)
            status_response.raise_for_status()
            status_data = status_response.json()
            status = status_data.get("status")

            if status == "succeeded":
                # Find image URL in all possible keys
                image_url = (
                    status_data.get("image_url")
                    or status_data.get("result", {}).get("image_url")
                    or (status_data.get("output")[0] if isinstance(status_data.get("output"), list) else status_data.get("output"))
                )
                if not image_url:
                    return None, f"Error: No image URL found in response: {status_data}"

                # Download the generated image
                img_resp = requests.get(image_url, timeout=30)
                img_resp.raise_for_status()
                output_path = Path("output.png")
                with open(output_path, "wb") as f:
                    f.write(img_resp.content)

                return str(output_path), None

            elif status == "failed":
                return None, "Error: Generation failed."

            time.sleep(5)

        return None, "Error: Timed out."

    except requests.exceptions.RequestException as e:
        return None, f"Request error: {str(e)}"

# === Gradio Interface ===
def process(reference_image, prompt):
    image_path, error = generate_image(reference_image, prompt)
    if error:
        return None, None, error
    download_link = image_path if image_path else None
    return image_path, download_link, "✅ Image generated successfully!"

iface = gr.Interface(
    fn=process,
    inputs=[
        gr.Image(type="filepath", label="Upload Reference Face Image"),
        gr.Textbox(lines=2, placeholder="Optional prompt", label="Prompt")
    ],
    outputs=[
        gr.Image(label="Generated Image"),
        gr.File(label="Download Image"),
        gr.Textbox(label="Status")
    ],
    title="Kling AI - Single Reference Face Generator",
    description="Upload a reference face image. The model will generate an image based on that face with 97% strength."
)

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