Update app.py

app.py

@@ -1,114 +1,186 @@
 import gradio as gr
 import requests
 import base64
-from pathlib import Path
-import jwt
+import os
 import time
+import jwt
 import logging
+from pathlib import Path
 
-# Set up logging for debugging
-logging.basicConfig(level=logging.DEBUG)
+# Set up logging
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Kling AI API configuration
+# ===== CONFIGURATION =====
 ACCESS_KEY_ID = "AFyHfnQATghFdCMyAG3gRPbNY4TNKFGB"
 ACCESS_KEY_SECRET = "TTepeLyBterLNM3brYPGmdndBnnyKJBA"
-API_URL = "https://api-singapore.klingai.com/v1/images/generations"
+BASE_URL = "https://api-singapore.klingai.com"
+CREATE_TASK_URL = f"{BASE_URL}/v1/images/generations"
 
+# ===== UTILITY FUNCTIONS =====
 def generate_jwt_token():
-    """Generate JWT token for Kling AI API authentication."""
+    """Generate JWT token for API authentication"""
     payload = {
         "iss": ACCESS_KEY_ID,
-        "exp": int(time.time()) + 1800,
-        "nbf": int(time.time()) - 5
+        "exp": int(time.time()) + 1800,  # Expires in 30 mins
+        "nbf": int(time.time()) - 5      # Not before 5 seconds ago
     }
-    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."
+    return jwt.encode(payload, ACCESS_KEY_SECRET, algorithm="HS256")
 
+def validate_image(image_path):
+    """Check image meets API requirements"""
     try:
-        with open(image, "rb") as img_file:
-            image_base64 = base64.b64encode(img_file.read()).decode("utf-8")
+        img_size = os.path.getsize(image_path) / 1024 / 1024  # MB
+        if img_size > 10:
+            return False, "Image too large (max 10MB)"
+        
+        # Add actual dimension check if possible (requires PIL)
+        return True, ""
     except Exception as e:
-        return None, f"Error: Failed to process image. Details: {str(e)}"
+        return False, f"Image validation error: {str(e)}"
 
+# ===== API FUNCTIONS =====
+def create_image_task(image_base64, prompt):
+    """Create image generation task"""
     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,
+        "model_name": "kling-v2",  # Best for image-to-image
+        "prompt": prompt,
+        "image": image_base64,
+        "resolution": "2k",
+        "n": 1,
         "aspect_ratio": "1:1"
     }
+    
+    try:
+        response = requests.post(CREATE_TASK_URL, json=payload, headers=headers)
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        logger.error(f"API request failed: {str(e)}")
+        return None
 
+def get_task_result(task_id):
+    """Retrieve task results"""
+    headers = {"Authorization": f"Bearer {generate_jwt_token()}"}
+    task_url = f"{BASE_URL}/v1/images/generations/{task_id}"
+    
     try:
-        logger.debug(f"Sending POST to {API_URL} with payload")
-        response = requests.post(API_URL, json=payload, headers=headers, timeout=30)
+        response = requests.get(task_url, headers=headers)
         response.raise_for_status()
-        data = response.json()
-        logger.debug(f"API response: {data}")
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        logger.error(f"Task status check failed: {str(e)}")
+        return None
 
-        task_id = data.get("data", {}).get("task_id")
-        if not task_id:
-            return None, "Error: No task ID returned."
+# ===== MAIN PROCESSING =====
+def generate_image(image_path, prompt=""):
+    """Handle end-to-end image generation"""
+    # Validate input image
+    is_valid, error_msg = validate_image(image_path)
+    if not is_valid:
+        return None, error_msg
 
-        # 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")
+    # Prepare image data
+    try:
+        with open(image_path, "rb") as img_file:
+            image_base64 = base64.b64encode(img_file.read()).decode("utf-8")
+    except Exception as e:
+        return None, f"Image processing error: {str(e)}"
+
+    # Create generation task
+    task_response = create_image_task(image_base64, prompt or "Transform into a vibrant cartoon style")
+    if not task_response or task_response.get("code") != 0:
+        return None, "Failed to create task"
+
+    task_id = task_response["data"]["task_id"]
+    logger.info(f"Created task: {task_id}")
+
+    # Poll for results (max 10 mins)
+    for _ in range(60):
+        task_data = get_task_result(task_id)
+        if not task_data:
+            time.sleep(5)
+            continue
+            
+        status = task_data["data"]["task_status"]
+        
+        if status == "succeed":
+            image_url = task_data["data"]["task_result"]["images"][0]["url"]
+            try:
+                img_data = requests.get(image_url).content
+                output_path = Path(f"/tmp/kling_output_{task_id}.png")
                 with open(output_path, "wb") as f:
-                    f.write(img_resp.content)
+                    f.write(img_data)
                 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)}"
+            except Exception as e:
+                return None, f"Failed to save image: {str(e)}"
+                
+        elif status in ("failed", "canceled"):
+            return None, f"Task failed: {task_data.get('task_status_msg', 'Unknown error')}"
+            
+        time.sleep(10)
+    
+    return None, "Task timed out after 10 minutes"
 
+# ===== GRADIO INTERFACE =====
 def chatbot_interface(image, prompt):
+    if not image:
+        return None, None, "Please upload an image first"
+    
     output_path, error = generate_image(image, prompt)
     if error:
         return None, None, error
-    return output_path, output_path, None
+    
+    return output_path, output_path, "Generation successful!"
 
-# Define Gradio interface
-with gr.Blocks() as iface:
-    gr.Markdown("# Kling AI Image-to-Image Generator")
+with gr.Blocks(title="Kling AI Image Transformer") as app:
+    gr.Markdown("# 🎨 Kling AI Image-to-Image Generator")
+    gr.Markdown("Transform images using Kling AI's Kolors technology")
+    
     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")
+            gr.Markdown("## Input Settings")
+            image_input = gr.Image(
+                type="filepath", 
+                label="Upload Image",
+                sources=["upload"],
+            )
+            prompt_input = gr.Textbox(
+                lines=2,
+                label="Transformation Prompt",
+                placeholder="Describe how you want to transform the image (e.g. 'vibrant watercolor painting')"
+            )
+            generate_btn = gr.Button("Generate", variant="primary")
+            
+            gr.Markdown("### Requirements")
+            gr.Markdown("""
+            - Max image size: 10MB
+            - Min dimensions: 300x300px
+            - Supported formats: JPG, PNG
+            - Aspect ratio: Between 1:2.5 and 2.5:1
+            """)
+            
         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(
+            gr.Markdown("## Output")
+            output_image = gr.Image(label="Generated Image", interactive=False)
+            output_file = gr.File(label="Download Result", file_types=["image/png"])
+            status_output = gr.Textbox(label="Status", interactive=False)
+    
+    generate_btn.click(
         fn=chatbot_interface,
         inputs=[image_input, prompt_input],
-        outputs=[output_image, output_file, error_message]
+        outputs=[output_image, output_file, status_output]
     )
 
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    app.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True
+    )