Update app.py

app.py

@@ -5,212 +5,110 @@ from pathlib import Path
 import jwt
 import time
 import logging
-import os
 
 # Set up logging for debugging
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
-# Kling AI API configuration (new keys)
+# Kling AI API configuration
 ACCESS_KEY_ID = "AFyHfnQATghFdCMyAG3gRPbNY4TNKFGB"
 ACCESS_KEY_SECRET = "TTepeLyBterLNM3brYPGmdndBnnyKJBA"
-API_URLS = [
-    "https://api-singapore.klingai.com/v1/images/image2image",  # Primary endpoint
-    "https://api-singapore.klingai.com/v1/images/kolors",  # Fallback for Kolors 2.1
-    "https://api-singapore.klingai.com/v1/images/generations"  # Fallback
-]
+API_URL = "https://api-singapore.klingai.com/v1/images/generations"
 
 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)
-    logger.debug(f"Generated JWT token: {token}")
-    return token
+    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 single-reference face generation.
-    
-    Args:
-        image: Uploaded image file (face image)
-        prompt (str): Optional text prompt
-    
-    Returns:
-        tuple: (Path to generated image or None, error message or None)
-    """
+    """Call Kling AI API for image-to-image generation."""
     if not image:
-        logger.error("No image uploaded")
-        return None, "Error: Please upload a valid face image (PNG/JPEG, <10 MB, ≥512x512 pixels)."
-    
+        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")
-        logger.debug("Image converted to base64 successfully")
     except Exception as e:
-        logger.error(f"Failed to process image: {str(e)}")
-        return None, f"Error: Failed to process image. Ensure it’s a valid PNG/JPEG. Details: {str(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_name": "kolors-v2.1",  # Kolors 2.1 for image-to-image
-        "image": image_base64,
+        "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",
-        "n": 1,
+        "num_outputs": 1,
         "aspect_ratio": "1:1"
     }
-    
-    for api_url in API_URLS:
-        try:
-            logger.debug(f"Sending POST request to {api_url} with payload: {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("task_id") or data.get("taskId") or data.get("id") or data.get("data", {}).get("task_id")
-            if not task_id:
-                logger.error("No task ID returned in API response")
-                return None, "Error: No task ID returned. Check API endpoint at https://klingai.com/global/dev/model/image."
-            
-            status_urls = [
-                f"https://api-singapore.klingai.com/v1/tasks/{task_id}",
-                f"https://api-singapore.klingai.com/api/v1/task/{task_id}"
-            ]
-            for status_url in status_urls:
-                for _ in range(60):
-                    logger.debug(f"Polling status at {status_url}")
-                    status_response = requests.get(status_url, headers=headers, timeout=30)
-                    if status_response.status_code == 404 and status_url != status_urls[-1]:
-                        logger.debug(f"Status endpoint {status_url} failed, trying next")
-                        continue
-                    status_response.raise_for_status()
-                    status_data = status_response.json()
-                    logger.debug(f"Status response: {status_data}")
-                    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("result", {}).get("image_url") or 
-                                   status_data.get("output") or 
-                                   status_data.get("data", {}).get("image") or
-                                   status_data.get("data", {}).get("images", [{}])[0].get("url"))
-                        if not image_url:
-                            logger.error("No image URL in API response")
-                            return None, "Error: No image URL in API response. Check https://klingai.com/global/dev/model/image."
-                        logger.debug(f"Downloading image from {image_url}")
-                        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)
-                        if not output_path.exists():
-                            logger.error("Output image file not created")
-                            return None, "Error: Failed to save output image."
-                        logger.debug(f"Image saved to {output_path}")
-                        return str(output_path), None
-                    elif status == "failed":
-                        logger.error("Image generation failed")
-                        return None, "Error: Image generation failed. Ensure the image contains a clear face and avoid NSFW content."
-                    elif status == "processing" and status_data.get("progress", 0) >= 0.99:
-                        logger.error("Generation stuck at 99%")
-                        return None, "Error: Generation stuck at 99%. Check account credits at https://app.klingai.com."
-                    time.sleep(5)
-                
-                logger.error("Image generation timed out")
-                return None, "Error: Image generation timed out. Try at 2 AM IST or check account status."
-            
-            return None, f"Error: All status endpoints failed. Check https://klingai.com/global/dev/model/image."
-        
-        except requests.exceptions.HTTPError as e:
-            status_code = e.response.status_code if e.response else None
-            error_data = e.response.json() if e.response and e.response.text else {}
-            service_code = error_data.get("code", 0)
-            error_message = error_data.get("message", str(e))
-            logger.error(f"HTTP error: {status_code}, Service Code: {service_code}, Details: {error_message}")
-            
-            if error_message.startswith("Create task failed"):
-                logger.error("Create task failed error detected")
-                return None, "Error: Create task failed. Try at 2 AM IST, verify account at https://app.klingai.com, or join Kling AI Discord."
-            elif status_code == 404 and api_url != API_URLS[-1]:
-                logger.debug(f"Endpoint {api_url} failed with 404, trying next")
-                continue
-            elif status_code == 404:
-                return None, f"Error: Endpoint not found (404). Tried {', '.join(API_URLS)}. Verify endpoint at https://klingai.com/global/dev/model/image."
-            elif status_code == 401:
-                if service_code in (1000, 1001, 1002, 1003, 1004):
-                    return None, f"Error: Authentication failed ({service_code}). Verify API keys or token settings."
-            elif status_code == 429:
-                if service_code in (1100, 1101, 1102, 1302, 1303, 1304):
-                    return None, f"Error: Rate limit or account issue ({service_code}). Try at 2 AM IST or check https://app.klingai.com."
-            elif status_code == 403 and service_code == 1103:
-                return None, "Error: Unauthorized access (1103). Verify account permissions at https://app.klingai.com."
-            elif status_code == 400:
-                if service_code in (1200, 1201, 1300, 1301):
-                    return None, f"Error: Invalid parameters ({service_code}). Use PNG/JPEG (<10 MB, ≥512x512) with a clear face."
-            elif status_code in (500, 503, 504):
-                if service_code in (5000, 5001, 5002):
-                    return None, f"Error: Server issue ({service_code}). Try at 2 AM IST or contact support."
-            return None, f"Error: API request failed. HTTP {status_code}, Service Code {service_code}. Details: {error_message}"
-        except requests.exceptions.RequestException as e:
-            logger.error(f"Network error: {str(e)}")
-            return None, f"Error: Network issue. Ensure stable internet, disable VPN. Details: {str(e)}"
-    
-    return None, f"Error: All API endpoints failed. Verify endpoint at https://klingai.com/global/dev/model/image."
+
+    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):
-    """
-    Gradio interface for single-reference face generation.
-    
-    Args:
-        image: Uploaded image file (containing a face)
-        prompt (str): Optional text prompt
-    
-    Returns:
-        tuple: (Output image path for display, output image path for download, error message)
-    """
     output_path, error = generate_image(image, prompt)
     if error:
-        logger.debug(f"Error returned: {error}")
         return None, None, error
-    logger.debug(f"Returning image path for display and download: {output_path}")
     return output_path, output_path, None
 
 # Define Gradio interface
 with gr.Blocks() as iface:
-    gr.Markdown(
-        """
-        # Kling AI Single-Reference Face Generator
-        Upload a PNG/JPEG image (<10 MB, ≥512x512 pixels) with a clear face to generate a transformed image using Kling AI Kolors 2.1 (strength 0.97). Avoid NSFW content.
-        If errors persist, verify endpoint at https://klingai.com/global/dev/model/image or contact support at https://app.klingai.com or Kling AI Discord.
-        """
-    )
+    gr.Markdown("# Kling AI Image-to-Image Generator")
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(type="filepath", label="Upload Face Image (PNG/JPEG, <10 MB, ≥512x512)")
-            prompt_input = gr.Textbox(lines=2, placeholder="Enter an optional prompt (e.g., 'Turn this face into a cartoon')", label="Prompt")
+            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 Face Image")
+            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]
     )
 
-# Launch the interface
 if __name__ == "__main__":
-    logger.debug("Launching Gradio interface")
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    iface.launch(server_name="0.0.0.0", server_port=7860)