Update app.py

app.py

@@ -1,274 +1,230 @@
-# app.py
-import os
-import time
-import base64
-import json
-from pathlib import Path
-
 import gradio as gr
 import requests
+import base64
+from pathlib import Path
 import jwt
-from PIL import Image
+import time
+import logging
+import os
 
-# ──────────────────────────────────────────────────────────────────────────────
-# CONFIG — set your keys as HF Space secrets or env vars for safety.
-# (Falls back to the keys you shared.)
-# ──────────────────────────────────────────────────────────────────────────────
-ACCESS_KEY_ID = os.getenv("KLING_ACCESS_KEY_ID", "AGBGmadNd9hakFYfahytyQQJtN8CJmDJ")
-ACCESS_KEY_SECRET = os.getenv("KLING_ACCESS_KEY_SECRET", "dp3pAe4PpdmnAHCAPgEd3PyLmBQrkMde")
+# Set up logging for debugging
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
 
-API_BASE = "https://api.klingai.com"
-ENDPOINT_KOLORS = f"{API_BASE}/v1/images/kolors"           # face/subject reference modes (image-to-image)
-ENDPOINT_GENERATIONS = f"{API_BASE}/v1/images/generations"  # listing (used as a fallback poller)
-ENDPOINT_TASK = lambda tid: f"{API_BASE}/v1/tasks/{tid}"    # primary poller
+# Kling AI API configuration (keys hardcoded as requested)
+ACCESS_KEY_ID = "AGBGmadNd9hakFYfahytyQQJtN8CJmDJ"
+ACCESS_KEY_SECRET = "dp3pAe4PpdmnAHCAPgEd3PyLmBQrkMde"
+API_URL = "https://api.klingai.com/v1/images/image2image"  # Correct endpoint
 
-# ──────────────────────────────────────────────────────────────────────────────
-# AUTH — Kling uses JWT: iss / exp / nbf with HS256 (no "access_key" field)
-# ──────────────────────────────────────────────────────────────────────────────
-def make_jwt() -> str:
-    headers = {"alg": "HS256", "typ": "JWT"}
-    now = int(time.time())
+def generate_jwt_token():
+    """Generate JWT token for Kling AI API authentication."""
+    headers = {
+        "alg": "HS256",
+        "typ": "JWT"
+    }
     payload = {
         "iss": ACCESS_KEY_ID,
-        "exp": now + 1800,  # 30 minutes
-        "nbf": now - 5,     # start now (minus small skew)
+        "exp": int(time.time()) + 1800,  # Token expires in 30 minutes
+        "nbf": int(time.time()) - 5  # Effective 5 seconds before current time
     }
-    return jwt.encode(payload, ACCESS_KEY_SECRET, algorithm="HS256", headers=headers)
-
-# ──────────────────────────────────────────────────────────────────────────────
-# HELPERS
-# ──────────────────────────────────────────────────────────────────────────────
-def ensure_image_ok(img_path: str):
-    with Image.open(img_path) as im:
-        im.verify()  # quick integrity check
+    token = jwt.encode(payload, ACCESS_KEY_SECRET, headers=headers)
+    logger.debug(f"Generated JWT token: {token}")
+    return token
 
-def b64_data_uri(img_path: str) -> str:
-    mime = "image/png" if img_path.lower().endswith(".png") else "image/jpeg"
-    with open(img_path, "rb") as f:
-        b = base64.b64encode(f.read()).decode("utf-8")
-    return f"data:{mime};base64,{b}"
-
-def extract_task_id(resp_json: dict) -> str | None:
-    # Common shapes seen in the wild
-    if not resp_json:
-        return None
-    for key in ("task_id", "taskId", "id"):
-        if key in resp_json:
-            return str(resp_json[key])
-    data = resp_json.get("data") or {}
-    for key in ("task_id", "taskId", "id"):
-        if key in data:
-            return str(data[key])
-    # Sometimes nested deeper (e.g., {"task": {"id": ...}})
-    task = resp_json.get("task") or data.get("task") or {}
-    if "id" in task:
-        return str(task["id"])
-    return None
-
-def extract_image_urls(resp_json: dict) -> list[str]:
-    if not resp_json:
-        return []
-    data = resp_json.get("data") or {}
-    # Typical: data.task_result.images = [{url: "..."}]
-    task_result = data.get("task_result") or {}
-    images = task_result.get("images") or []
-    urls = [img.get("url") for img in images if isinstance(img, dict) and img.get("url")]
-    if urls:
-        return urls
-    # Some variants: output, image_url, result.image_url
-    for k in ("output", "image_url"):
-        if k in resp_json and isinstance(resp_json[k], str):
-            return [resp_json[k]]
-    result = resp_json.get("result") or {}
-    if isinstance(result, dict) and result.get("image_url"):
-        return [result["image_url"]]
-    # Works array pattern
-    works = resp_json.get("works") or data.get("works") or []
-    urls = []
-    for w in works:
-        if isinstance(w, dict):
-            u = w.get("url") or w.get("imageUrl")
-            if u:
-                urls.append(u)
-    return urls
-
-def download_to_file(url: str, out_path: Path) -> Path:
-    r = requests.get(url, timeout=60)
-    r.raise_for_status()
-    out_path.parent.mkdir(parents=True, exist_ok=True)
-    with open(out_path, "wb") as f:
-        f.write(r.content)
-    return out_path
-
-def poll_for_result(task_id: str, headers: dict, timeout_s: int = 300, interval_s: float = 3.0) -> dict:
-    """Poll task endpoint first; fallback to listing."""
-    deadline = time.time() + timeout_s
-    last_error = None
-
-    while time.time() < deadline:
-        try:
-            # Preferred: direct task status
-            r = requests.get(ENDPOINT_TASK(task_id), headers=headers, timeout=30)
-            if r.status_code == 200:
-                j = r.json()
-                # Either "status_name":"succeed" or "data.task_status":"succeed"
-                status_name = (j.get("status_name")
-                               or (j.get("data") or {}).get("task_status")
-                               or (j.get("task") or {}).get("status_name"))
-                if isinstance(status_name, dict):
-                    # Some SDKs wrap status as enum-like
-                    status_name = status_name.get("value")
-                if status_name in ("succeed", "succeeded", "success",  "SUCCEED"):
-                    return j
-                if status_name in ("failed", "FAIL", "failed_with_error"):
-                    return j
-            elif r.status_code in (401, 403, 404):
-                last_error = r.text
-            # Fallback: scan generations list
-            r2 = requests.get(ENDPOINT_GENERATIONS, headers=headers, params={"pageSize": 200}, timeout=30)
-            if r2.status_code == 200:
-                j2 = r2.json()
-                for item in (j2.get("data") or []):
-                    if str(item.get("task_id")) == str(task_id):
-                        status = item.get("task_status")
-                        if status in ("succeed", "succeeded", "success"):
-                            return item
-                        if status in ("failed",):
-                            return item
-        except requests.RequestException as e:
-            last_error = str(e)
-        time.sleep(interval_s)
-    raise TimeoutError(f"Polling timed out for task_id {task_id}. Last error: {last_error or 'n/a'}")
-
-# ──────────────────────────────────────────────────────────────────────────────
-# CORE CALL — Kolors face reference (single reference, faceStrength=97)
-# ──────────────────────────────────────────────────────────────────────────────
-def kling_face_reference(image_path: str, prompt: str, face_strength: int = 97, aspect_ratio: str = "1:1") -> tuple[str, str]:
+def generate_image(image, prompt=""):
     """
-    Returns (display_image_path, download_file_path)
+    Call Kling AI API for single-reference face generation.
+    
+    Args:
+        image: Uploaded image file (from Gradio, face image)
+        prompt (str): Optional text prompt to guide transformation
+    
+    Returns:
+        tuple: (Path to generated image or None, error message or None)
     """
-    if not image_path:
-        raise gr.Error("Please upload a face/reference image.")
-    ensure_image_ok(image_path)
-
-    token = make_jwt()
-    headers_json = {
-        "Authorization": f"Bearer {token}",
-        "Content-Type": "application/json",
-    }
-    headers_multipart = {
-        "Authorization": f"Bearer {token}",
-    }
-
-    # First try: multipart/form-data (send file as `imageReference`)
-    data_multipart = {
-        "prompt": (None, prompt),
-        "reference": (None, "face"),
-        "faceStrength": (None, str(max(1, min(100, int(face_strength))))),
-        "faceNo": (None, "1"),           # single face reference
-        "imageCount": (None, "1"),
-        "aspect_ratio": (None, aspect_ratio),
+    if not image:
+        logger.error("No image uploaded")
+        return None, "Error: Please upload a valid face image (PNG/JPEG, <10 MB, ≥512x512 pixels)."
+    
+    # Convert image to base64
+    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)}"
+    
+    headers = {
+        "Authorization": f"Bearer {generate_jwt_token()}",
+        "Content-Type": "application/json"
     }
-    files = {
-        "imageReference": (os.path.basename(image_path), open(image_path, "rb"),
-                           "image/png" if image_path.lower().endswith(".png") else "image/jpeg")
+    
+    payload = {
+        "model_name": "kling-v2-1",  # V2.1 model for image-to-image
+        "task_type": "image_generation",  # Adapted from video task structure
+        "input": {
+            "image": image_base64,
+            "prompt": prompt or "Transform the face into a cartoon style while preserving identity",
+            "strength": 0.97,  # High reference strength
+            "output_format": "png"
+        }
     }
-
-    # Attempt 1 — multipart
+    
     try:
-        resp = requests.post(ENDPOINT_KOLORS, headers=headers_multipart, files=files, data=data_multipart, timeout=60)
-        if resp.status_code == 200:
-            j = resp.json()
-        else:
-            # Read JSON anyway if possible
-            try:
-                j = resp.json()
-            except Exception:
-                j = {"code": resp.status_code, "message": resp.text}
-    finally:
-        # Close file handle if opened
-        try:
-            files["imageReference"][1].close()
-        except Exception:
-            pass
-
-    task_id = extract_task_id(j)
-
-    # If Kolors rejected multipart or no task_id, try JSON with data URI
-    if not task_id:
-        payload = {
-            "prompt": prompt,
-            "reference": "face",
-            "faceStrength": max(1, min(100, int(face_strength))),
-            "faceNo": 1,
-            "imageCount": 1,
-            "aspect_ratio": aspect_ratio,
-            "imageReference": b64_data_uri(image_path),
-        }
-        resp2 = requests.post(ENDPOINT_KOLORS, headers=headers_json, json=payload, timeout=60)
-        try:
-            j = resp2.json()
-        except Exception:
-            j = {"code": resp2.status_code, "message": resp2.text}
-        task_id = extract_task_id(j)
-
-    if not task_id:
-        code = j.get("code") or j.get("service_code") or "?"
-        msg = j.get("message") or j.get("error") or f"HTTP {resp.status_code if 'resp' in locals() else '?'}"
-        raise gr.Error(f"Create task failed. Code: {code}. Message: {msg}")
-
-    # Poll
-    result_json = poll_for_result(task_id, headers=headers_json, timeout_s=420, interval_s=3.0)
-
-    # Gather image URLs
-    urls = extract_image_urls(result_json)
-    if not urls:
-        # Some APIs return the latest object on /v1/images/generations with same task_id
-        try:
-            listing = requests.get(ENDPOINT_GENERATIONS, headers=headers_json, params={"pageSize": 200}, timeout=30).json()
-            for item in (listing.get("data") or []):
-                if str(item.get("task_id")) == str(task_id):
-                    urls = extract_image_urls(item)
-                    if urls:
-                        break
-        except Exception:
-            pass
-
-    if not urls:
-        raise gr.Error(f"Task {task_id} succeeded but no image URL found in response.")
-
-    # Download first image
-    out_dir = Path("outputs")
-    out_dir.mkdir(parents=True, exist_ok=True)
-    out_path = out_dir / f"kling_face_{task_id}.png"
-    download_to_file(urls[0], out_path)
-
-    # Return same path for preview and download
-    return str(out_path), str(out_path)
-
-# ──────────────────────────────────────────────────────────────────────────────
-# GRADIO UI
-# ──────────────────────────────────────────────────────────────────────────────
-with gr.Blocks(title="Kling AI — Image to Image (Face Reference)") as demo:
-    gr.Markdown("### Kling AI — Image-to-Image (Single Face Reference)\nUpload a face image and a prompt. Strength defaults to 97.")
-    with gr.Row():
-        in_image = gr.Image(type="filepath", label="Reference Face Image (PNG/JPG)")
-        in_prompt = gr.Textbox(label="Prompt", placeholder="e.g., Ultra-detailed portrait, soft light, studio background", lines=2)
+        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")
+        if not task_id:
+            logger.error("No task ID returned in API response")
+            return None, "Error: No task ID returned. Check API endpoint or keys at https://app.klingai.com/global/dev."
+        
+        # Poll for task completion
+        status_url = f"https://api.klingai.com/v1/tasks/{task_id}"
+        for _ in range(60):  # Poll for up to 5 minutes
+            logger.debug(f"Polling status at {status_url}")
+            status_response = requests.get(status_url, headers=headers, timeout=30)
+            status_response.raise_for_status()
+            status_data = status_response.json()
+            logger.debug(f"Status response: {status_data}")
+            status = status_data.get("status")
+            if status == "succeeded":
+                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")
+                if not image_url:
+                    logger.error("No image URL in API response")
+                    return None, "Error: No image URL in API response. Check API documentation at https://app.klingai.com/global/dev."
+                # Download the 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")  # Use /tmp for Spaces compatibility
+                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 or sensitive 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 or upgrade to a paid plan at https://app.klingai.com."
+            time.sleep(5)
+        
+        logger.error("Image generation timed out")
+        return None, "Error: Image generation timed out. Try during off-peak hours (e.g., 2 AM IST) or check account status."
+    
+    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 during off-peak hours (e.g., 2 AM IST), verify account status at https://app.klingai.com, or join Kling AI Discord for support."
+        elif status_code == 404:
+            if service_code == 1202:
+                return None, "Error: Invalid request method (1202). Check https://app.klingai.com/global/dev for correct endpoint."
+            elif service_code == 1203:
+                return None, "Error: Resource does not exist (1203). Verify model or endpoint."
+            return None, f"Error: Endpoint not found (404). Verify https://api.klingai.com/v1/images/image2image at https://app.klingai.com/global/dev."
+        elif status_code == 401:
+            if service_code == 1000:
+                return None, "Error: Authentication failed (1000). Verify API keys are correct."
+            elif service_code == 1001:
+                return None, "Error: Authorization empty (1001). Ensure JWT token is included."
+            elif service_code == 1002:
+                return None, "Error: Authorization invalid (1002). Check key format."
+            elif service_code == 1003:
+                return None, "Error: Authorization not yet valid (1003). Adjust token start time (nbf)."
+            elif service_code == 1004:
+                return None, "Error: Authorization expired (1004). Generate a new token."
+        elif status_code == 429:
+            if service_code == 1100:
+                return None, "Error: Account exception (1100). Verify account configuration at https://app.klingai.com."
+            elif service_code == 1101:
+                return None, "Error: Account in arrears (1101). Recharge your account."
+            elif service_code == 1102:
+                return None, "Error: Resource pack depleted or expired (1102). Purchase additional resources."
+            elif service_code in (1302, 1303):
+                return None, "Error: Rate limit exceeded (1302/1303). Try during off-peak hours (e.g., 2 AM IST) or contact support."
+            elif service_code == 1304:
+                return None, "Error: IP whitelisting issue (1304). Contact Kling AI support."
+        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 == 1200:
+                return None, "Error: Invalid request parameters (1200). Check payload format."
+            elif service_code == 1201:
+                return None, "Error: Invalid parameters (1201). Use PNG/JPEG (<10 MB, ≥512x512) with a clear face."
+            elif service_code == 1300:
+                return None, "Error: Platform policy triggered (1300). Check input content."
+            elif service_code == 1301:
+                return None, "Error: Content security policy triggered (1301). Ensure the image contains a clear face and avoid NSFW content."
+        elif status_code in (500, 503, 504):
+            if service_code == 5000:
+                return None, "Error: Server internal error (5000). Try again later."
+            elif service_code == 5001:
+                return None, "Error: Server unavailable due to maintenance (5001). Try again later."
+            elif service_code == 5002:
+                return None, "Error: Server timeout (5002). Try during off-peak hours (e.g., 2 AM IST)."
+        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, and try again. 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 V2.1 API (reference strength 0.97). Avoid NSFW or sensitive content.
+        If errors persist, check account status at https://app.klingai.com or join Kling AI Discord for support.
+        """
+    )
     with gr.Row():
-        in_strength = gr.Slider(1, 100, value=97, step=1, label="Face Reference Strength")
-        in_aspect = gr.Dropdown(choices=["1:1", "3:4", "4:3", "2:3", "3:2", "16:9", "9:16", "21:9"], value="1:1", label="Aspect Ratio")
-    btn = gr.Button("Generate", variant="primary")
-
-    out_img = gr.Image(label="Generated Image", show_download_button=False)
-    out_file = gr.File(label="Download Image")
-
-    def run(image, prompt, strength, aspect):
-        if not prompt or not prompt.strip():
-            raise gr.Error("Please enter a prompt.")
-        return kling_face_reference(image, prompt.strip(), int(strength), aspect)
-
-    btn.click(fn=run, inputs=[in_image, in_prompt, in_strength, in_aspect], outputs=[out_img, out_file])
-
+        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")
+            generate_button = gr.Button("Generate")
+        with gr.Column():
+            output_image = gr.Image(label="Generated Face 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__":
-    # On HF Spaces, just `python app.py` is enough — no need to set host/port.
-    demo.launch()
+    logger.debug("Launching Gradio interface")
+    iface.launch(server_name="0.0.0.0", server_port=7860)