Update app.py

app.py

@@ -7,34 +7,45 @@ import jwt
 import logging
 from pathlib import Path
 
-# Set up logging
+# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# ===== CONFIGURATION =====
+# ===== API CONFIGURATION =====
 ACCESS_KEY_ID = "AFyHfnQATghFdCMyAG3gRPbNY4TNKFGB"
 ACCESS_KEY_SECRET = "TTepeLyBterLNM3brYPGmdndBnnyKJBA"
-BASE_URL = "https://api-singapore.klingai.com"
-CREATE_TASK_URL = f"{BASE_URL}/v1/images/generations"
+API_BASE_URL = "https://api-singapore.klingai.com"
+CREATE_TASK_ENDPOINT = f"{API_BASE_URL}/v1/images/generations"
 
-# ===== UTILITY FUNCTIONS =====
+# ===== AUTHENTICATION =====
 def generate_jwt_token():
     """Generate JWT token for API authentication"""
     payload = {
         "iss": ACCESS_KEY_ID,
-        "exp": int(time.time()) + 1800,  # Expires in 30 mins
-        "nbf": int(time.time()) - 5      # Not before 5 seconds ago
+        "exp": int(time.time()) + 1800,  # 30 minutes expiration
+        "nbf": int(time.time()) - 5     # Not before 5 seconds ago
     }
-    return jwt.encode(payload, ACCESS_KEY_SECRET, algorithm="HS256")
+    try:
+        token = jwt.encode(payload, ACCESS_KEY_SECRET, algorithm="HS256")
+        return token if isinstance(token, str) else token.decode('utf-8')
+    except Exception as e:
+        logger.error(f"JWT generation failed: {str(e)}")
+        return None
 
+# ===== IMAGE VALIDATION =====
 def validate_image(image_path):
-    """Check image meets API requirements"""
+    """Check if image meets API requirements"""
     try:
-        img_size = os.path.getsize(image_path) / 1024 / 1024  # MB
-        if img_size > 10:
+        # Check file size
+        size_mb = os.path.getsize(image_path) / (1024 * 1024)
+        if size_mb > 10:
             return False, "Image too large (max 10MB)"
         
-        # Add actual dimension check if possible (requires PIL)
+        # Check file extension
+        ext = os.path.splitext(image_path)[1].lower()
+        if ext not in ['.jpg', '.jpeg', '.png']:
+            return False, "Invalid format (only JPG/PNG)"
+            
         return True, ""
     except Exception as e:
         return False, f"Image validation error: {str(e)}"
@@ -42,137 +53,156 @@ def validate_image(image_path):
 # ===== API FUNCTIONS =====
 def create_image_task(image_base64, prompt):
     """Create image generation task"""
+    token = generate_jwt_token()
+    if not token:
+        return None, "Authentication failed"
+    
     headers = {
-        "Authorization": f"Bearer {generate_jwt_token()}",
+        "Authorization": f"Bearer {token}",
         "Content-Type": "application/json"
     }
     
     payload = {
-        "model_name": "kling-v2",  # Best for image-to-image
-        "prompt": prompt,
+        "model_name": "kling-v2",  # Recommended model
+        "prompt": prompt or "Transform into a vibrant artistic style",
         "image": image_base64,
-        "resolution": "2k",
-        "n": 1,
-        "aspect_ratio": "1:1"
+        "resolution": "1k",  # Start with 1k for faster results
+        "aspect_ratio": "1:1",
+        "n": 1
     }
     
     try:
-        response = requests.post(CREATE_TASK_URL, json=payload, headers=headers)
+        response = requests.post(
+            CREATE_TASK_ENDPOINT,
+            json=payload,
+            headers=headers,
+            timeout=30
+        )
         response.raise_for_status()
-        return response.json()
+        return response.json(), None
     except requests.exceptions.RequestException as e:
         logger.error(f"API request failed: {str(e)}")
-        return None
+        if hasattr(e, 'response') and e.response:
+            logger.error(f"API response: {e.response.text}")
+        return None, f"API request failed: {str(e)}"
 
-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}"
+def check_task_status(task_id):
+    """Check task completion status"""
+    token = generate_jwt_token()
+    if not token:
+        return None, "Authentication failed"
+    
+    headers = {"Authorization": f"Bearer {token}"}
+    status_url = f"{API_BASE_URL}/v1/images/generations/{task_id}"
     
     try:
-        response = requests.get(task_url, headers=headers)
+        response = requests.get(status_url, headers=headers, timeout=30)
         response.raise_for_status()
-        return response.json()
+        return response.json(), None
     except requests.exceptions.RequestException as e:
-        logger.error(f"Task status check failed: {str(e)}")
-        return None
+        logger.error(f"Status check failed: {str(e)}")
+        return None, f"Status check failed: {str(e)}"
 
 # ===== MAIN PROCESSING =====
-def generate_image(image_path, prompt=""):
-    """Handle end-to-end image generation"""
-    # Validate input image
+def generate_image(image_path, prompt):
+    """Handle complete image generation workflow"""
+    # Validate image
     is_valid, error_msg = validate_image(image_path)
     if not is_valid:
         return None, error_msg
-
-    # Prepare image data
+    
+    # Read and encode image
     try:
         with open(image_path, "rb") as img_file:
-            image_base64 = base64.b64encode(img_file.read()).decode("utf-8")
+            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"
-
+        return None, f"Failed to process image: {str(e)}"
+    
+    # Create task
+    task_response, error = create_image_task(image_base64, prompt)
+    if error:
+        return None, error
+    
+    if task_response.get("code") != 0:
+        return None, f"API error: {task_response.get('message', 'Unknown error')}"
+    
     task_id = task_response["data"]["task_id"]
-    logger.info(f"Created task: {task_id}")
-
-    # Poll for results (max 10 mins)
+    logger.info(f"Task created: {task_id}")
+    
+    # Poll for results (max 10 minutes)
     for _ in range(60):
-        task_data = get_task_result(task_id)
-        if not task_data:
-            time.sleep(5)
-            continue
+        task_data, error = check_task_status(task_id)
+        if error:
+            return None, error
             
         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
+                response = requests.get(image_url, timeout=30)
+                response.raise_for_status()
                 output_path = Path(f"/tmp/kling_output_{task_id}.png")
                 with open(output_path, "wb") as f:
-                    f.write(img_data)
+                    f.write(response.content)
                 return str(output_path), None
             except Exception as e:
-                return None, f"Failed to save image: {str(e)}"
+                return None, f"Failed to download result: {str(e)}"
                 
         elif status in ("failed", "canceled"):
-            return None, f"Task failed: {task_data.get('task_status_msg', 'Unknown error')}"
+            error_msg = task_data["data"].get("task_status_msg", "Unknown error")
+            return None, f"Task failed: {error_msg}"
             
         time.sleep(10)
     
     return None, "Task timed out after 10 minutes"
 
 # ===== GRADIO INTERFACE =====
-def chatbot_interface(image, prompt):
+def process_interface(image, prompt):
     if not image:
         return None, None, "Please upload an image first"
     
     output_path, error = generate_image(image, prompt)
     if error:
+        logger.error(f"Generation failed: {error}")
         return None, None, error
     
-    return output_path, output_path, "Generation successful!"
+    return output_path, output_path, "Success! Image generated"
 
-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.Blocks(title="Kling AI Image Generator") as app:
+    gr.Markdown("## 🎨 Kling AI Image-to-Image Transformation")
     
     with gr.Row():
         with gr.Column():
-            gr.Markdown("## Input Settings")
+            gr.Markdown("### Input Settings")
             image_input = gr.Image(
-                type="filepath", 
+                type="filepath",
                 label="Upload Image",
-                sources=["upload"],
+                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')"
+                placeholder="Describe the desired style (e.g. '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
+            - **Image Requirements**:
+              - Max size: 10MB
+              - Formats: JPG, PNG
+              - Min dimensions: 300x300px
+              - Aspect ratio between 1:2.5 and 2.5:1
             """)
             
         with gr.Column():
-            gr.Markdown("## Output")
+            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,
+        fn=process_interface,
         inputs=[image_input, prompt_input],
         outputs=[output_image, output_file, status_output]
     )
@@ -181,6 +211,5 @@ if __name__ == "__main__":
     app.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False,
-        show_error=True
+        share=False
     )