Update app.py

app.py

@@ -6,10 +6,10 @@ import jwt
 import time
 import logging
 
-# ================== CONFIGURATION ==================
-ACCESS_KEY_ID = "AGBGmadNd9hakFYfahytyQQJtN8CJmDJ"
-ACCESS_KEY_SECRET = "dp3pAe4PpdmnAHCAPgEd3PyLmBQrkMde"
-API_URL = "https://api-singapore.klingai.com/v1/images/generations"  # New unified endpoint
+# ================== CONFIG ==================
+ACCESS_KEY_ID = "YOUR_ACCESS_KEY_ID"
+ACCESS_KEY_SECRET = "YOUR_ACCESS_KEY_SECRET"
+API_URL = "https://api-singapore.klingai.com/v1/images/generations"
 
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
@@ -17,15 +17,13 @@ logger = logging.getLogger(__name__)
 
 # ================== AUTH ==================
 def generate_jwt_token():
-    """Generate JWT token for Kling AI API authentication."""
     headers = {"alg": "HS256", "typ": "JWT"}
     payload = {
         "iss": ACCESS_KEY_ID,
         "exp": int(time.time()) + 1800,
         "nbf": int(time.time()) - 5
     }
-    token = jwt.encode(payload, ACCESS_KEY_SECRET, headers=headers)
-    return token
+    return jwt.encode(payload, ACCESS_KEY_SECRET, headers=headers)
 
 
 # ================== IMAGE GENERATION ==================
@@ -37,7 +35,7 @@ def generate_image(image, prompt=""):
         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)}"
+        return None, f"Error reading image: {e}"
 
     headers = {
         "Authorization": f"Bearer {generate_jwt_token()}",
@@ -45,88 +43,79 @@ def generate_image(image, prompt=""):
     }
 
     payload = {
-        "model_name": "kolors-v2.1",
-        "image_base64": [image_base64],  # New format: array of base64 images
-        "prompt": prompt or "Transform the face into a cartoon style while preserving identity",
-        "strength": 0.97,
-        "output_format": "png",
-        "n": 1,
-        "aspect_ratio": "1:1"
+        "model_name": "kling-v2-1",           # ✅ per docs
+        "prompt": prompt or "Transform face into artistic style",
+        "image": image_base64,                # ✅ base64 only, no prefix
+        "aspect_ratio": "1:1",
+        "n": 1,                                # number of images
+        "resolution": "1k"
     }
 
     try:
-        logger.debug(f"Sending POST request to {API_URL}")
+        logger.debug(f"POST {API_URL} payload={payload.keys()}")
         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("data", {}).get("task_id")
-        if not task_id:
-            return None, "Error: No task ID returned from API."
+        if data.get("code") != 0:
+            return None, f"API Error: {data.get('message')}"
 
-        # ================== POLLING ==================
-        status_url = f"https://api-singapore.klingai.com/v1/tasks/{task_id}"
-        for _ in range(60):
+        task_id = data["data"]["task_id"]
+        logger.debug(f"Task ID: {task_id}")
+
+        # Poll task status
+        status_url = f"{API_URL}/{task_id}"
+        for _ in range(60):  # up to ~5 minutes
             time.sleep(5)
-            status_response = requests.get(status_url, headers=headers, timeout=30)
-            if status_response.status_code == 404:
-                continue
-            status_data = status_response.json()
-            status = status_data.get("task_status") or status_data.get("status")
-
-            if status in ("succeeded", "completed"):
-                image_url = (
-                    status_data.get("image_url")
-                    or status_data.get("data", {}).get("image")
-                    or status_data.get("data", {}).get("images", [{}])[0].get("url")
-                )
-                if not image_url:
-                    return None, "Error: No image URL in API response."
-
-                image_response = requests.get(image_url, timeout=30)
-                image_response.raise_for_status()
-                output_path = Path("/tmp/output_image.png")
-                with open(output_path, "wb") as f:
-                    f.write(image_response.content)
-                return str(output_path), None
-
-            elif status == "failed":
-                return None, "Error: Image generation failed."
-
-        return None, "Error: Image generation timed out."
+            status_resp = requests.get(status_url, headers=headers, timeout=30)
+            status_data = status_resp.json()
+
+            if status_data.get("data", {}).get("task_status") == "succeed":
+                images = status_data["data"]["task_result"]["images"]
+                if images:
+                    img_url = images[0]["url"]
+                    img_resp = requests.get(img_url)
+                    img_resp.raise_for_status()
+                    out_path = Path("/tmp/output.png")
+                    with open(out_path, "wb") as f:
+                        f.write(img_resp.content)
+                    return str(out_path), None
+                else:
+                    return None, "No image returned."
+
+            elif status_data.get("data", {}).get("task_status") == "failed":
+                return None, f"Generation failed: {status_data['data'].get('task_status_msg')}"
+
+        return None, "Timed out waiting for result."
 
     except requests.exceptions.HTTPError as e:
-        return None, f"HTTP Error: {str(e)}"
-    except requests.exceptions.RequestException as e:
-        return None, f"Network Error: {str(e)}"
+        return None, f"HTTP Error: {e}"
+    except Exception as e:
+        return None, f"Unexpected Error: {e}"
 
 
-# ================== GRADIO INTERFACE ==================
+# ================== GRADIO UI ==================
 def chatbot_interface(image, prompt):
-    output_path, error = generate_image(image, prompt)
+    image_path, error = generate_image(image, prompt)
     if error:
         return None, None, error
-    return output_path, output_path, None
+    return image_path, image_path, None
 
 
 with gr.Blocks() as iface:
-    gr.Markdown("# Kling AI Image-to-Image Generator (Kolors v2.1)")
+    gr.Markdown("# Kling AI Image Generator — Official API Format")
 
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(type="filepath", label="Upload Face Image")
-            prompt_input = gr.Textbox(lines=2, placeholder="Enter prompt", label="Prompt")
-            generate_button = gr.Button("Generate")
+            img_in = gr.Image(type="filepath", label="Upload Image")
+            prompt_in = gr.Textbox(lines=2, placeholder="Enter prompt", label="Prompt")
+            btn = gr.Button("Generate")
         with gr.Column():
-            output_image = gr.Image(label="Generated Image")
-            output_file = gr.File(label="Download Image")
-            error_message = gr.Textbox(label="Status/Error", interactive=False)
-
-    generate_button.click(
-        fn=chatbot_interface,
-        inputs=[image_input, prompt_input],
-        outputs=[output_image, output_file, error_message]
-    )
+            img_out = gr.Image(label="Generated Image")
+            file_out = gr.File(label="Download Image")
+            err_out = gr.Textbox(label="Error / Status", interactive=False)
+
+    btn.click(fn=chatbot_interface, inputs=[img_in, prompt_in], outputs=[img_out, file_out, err_out])
 
 if __name__ == "__main__":
     iface.launch(server_name="0.0.0.0", server_port=7860)