File size: 4,265 Bytes
85a7fa6
a594839
1a8f2c7
85a7fa6
 
 
 
1a8f2c7
85a7fa6
 
 
 
c3f22c6
97f568a
 
c3f22c6
1a8f2c7
 
85a7fa6
3a257f2
fab14c7
1a8f2c7
 
3a257f2
c3f22c6
 
97f568a
85a7fa6
c3f22c6
85a7fa6
c3f22c6
 
85a7fa6
 
 
 
c3f22c6
 
1a8f2c7
 
 
a594839
c3f22c6
85a7fa6
c3f22c6
85a7fa6
c3f22c6
85a7fa6
97f568a
c3f22c6
85a7fa6
a594839
c3f22c6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
85a7fa6
 
 
 
 
 
85857de
85a7fa6
 
c3f22c6
85a7fa6
 
c3f22c6
 
85a7fa6
 
c3f22c6
85a7fa6
 
c3f22c6
85a7fa6
 
 
 
 
1a8f2c7
045423f
c3f22c6
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import gradio as gr
import requests
import base64
from pathlib import Path
import jwt
import time
import logging

# Set up logging for debugging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Kling AI API configuration
ACCESS_KEY_ID = "AFyHfnQATghFdCMyAG3gRPbNY4TNKFGB"
ACCESS_KEY_SECRET = "TTepeLyBterLNM3brYPGmdndBnnyKJBA"
API_URL = "https://api-singapore.klingai.com/v1/images/generations"

def generate_jwt_token():
    """Generate JWT token for Kling AI API authentication."""
    payload = {
        "iss": ACCESS_KEY_ID,
        "exp": int(time.time()) + 1800,
        "nbf": int(time.time()) - 5
    }
    token = jwt.encode(payload, ACCESS_KEY_SECRET, algorithm="HS256")
    return token if isinstance(token, str) else token.decode("utf-8")

def generate_image(image, prompt=""):
    """Call Kling AI API for image-to-image generation."""
    if not image:
        return None, "Error: Please upload a valid image."

    try:
        with open(image, "rb") as img_file:
            image_base64 = base64.b64encode(img_file.read()).decode("utf-8")
    except Exception as e:
        return None, f"Error: Failed to process image. Details: {str(e)}"

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

    payload = {
        "model": "kolors-v2.1",
        "prompt": prompt or "Transform the face into a cartoon style while preserving identity",
        "image_base64_list": [image_base64],
        "strength": 0.97,
        "output_format": "png",
        "num_outputs": 1,
        "aspect_ratio": "1:1"
    }

    try:
        logger.debug(f"Sending POST to {API_URL} with payload")
        response = requests.post(API_URL, json=payload, headers=headers, timeout=30)
        response.raise_for_status()
        data = response.json()
        logger.debug(f"API response: {data}")

        task_id = data.get("data", {}).get("task_id")
        if not task_id:
            return None, "Error: No task ID returned."

        # Poll task until completion
        status_url = f"https://api-singapore.klingai.com/v1/images/task/{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()
            logger.debug(f"Status: {status_data}")
            task_status = status_data.get("data", {}).get("task_status")
            if task_status == "succeed":
                image_url = status_data["data"]["images"][0]["url"]
                img_resp = requests.get(image_url, timeout=30)
                output_path = Path("/tmp/output_image.png")
                with open(output_path, "wb") as f:
                    f.write(img_resp.content)
                return str(output_path), None
            elif task_status in ("failed", "canceled"):
                return None, f"Error: Task failed. Details: {status_data}"
            time.sleep(5)
        return None, "Error: Task polling timed out."

    except requests.exceptions.RequestException as e:
        return None, f"Error: API request failed. Details: {str(e)}"

def chatbot_interface(image, prompt):
    output_path, error = generate_image(image, prompt)
    if error:
        return None, None, error
    return output_path, output_path, None

# Define Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# Kling AI Image-to-Image Generator")
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="Upload Image")
            prompt_input = gr.Textbox(lines=2, placeholder="Enter prompt", label="Prompt")
            generate_button = gr.Button("Generate")
        with gr.Column():
            output_image = gr.Image(label="Generated Image")
            output_file = gr.File(label="Download Generated Image")
            error_message = gr.Textbox(label="Status/Error Message", interactive=False)

    generate_button.click(
        fn=chatbot_interface,
        inputs=[image_input, prompt_input],
        outputs=[output_image, output_file, error_message]
    )

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