Update app.py

app.py

@@ -15,7 +15,7 @@ logger = logging.getLogger(__name__)
 ACCESS_KEY_ID = "AFyHfnQATghFdCMyAG3gRPbNY4TNKFGB"
 ACCESS_KEY_SECRET = "TTepeLyBterLNM3brYPGmdndBnnyKJBA"
 API_BASE_URL = "https://api-singapore.klingai.com"
-CREATE_TASK_ENDPOINT = f"{API_BASE_URL}/v1/images/generations"  # SINGLE image endpoint
+CREATE_TASK_ENDPOINT = f"{API_BASE_URL}/v1/images/multi-image2image"
 
 # ===== AUTHENTICATION =====
 def generate_jwt_token():
@@ -37,146 +37,177 @@ def prepare_image_base64(image_path):
         logger.error(f"Image processing failed: {str(e)}")
         return None
 
-def validate_face_image(image_path):
-    """Validate the image meets face transformation requirements"""
+def validate_image(image_path):
+    """Validate image meets API requirements"""
     try:
-        # Check file exists
-        if not os.path.exists(image_path):
-            return False, "Image file not found"
-            
-        # Check file size (max 10MB)
-        file_size = os.path.getsize(image_path) / (1024 * 1024)
-        if file_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)"
-            
+        
+        # Check dimensions (basic check - should use PIL for actual dimensions)
         return True, ""
     except Exception as e:
-        return False, f"Validation error: {str(e)}"
+        return False, f"Image validation error: {str(e)}"
 
 # ===== API FUNCTIONS =====
-def create_face_task(image_base64, prompt):
-    """Create face transformation task with 97% fidelity"""
+def create_multi_image_task(subject_images, prompt):
+    """Create multi-image generation task"""
     headers = {
         "Authorization": f"Bearer {generate_jwt_token()}",
         "Content-Type": "application/json"
     }
     
+    # Prepare subject images list
+    subject_image_list = []
+    for img_path in subject_images:
+        if img_path:  # Skip empty/None images
+            base64_img = prepare_image_base64(img_path)
+            if base64_img:
+                subject_image_list.append({"subject_image": base64_img})
+    
+    if len(subject_image_list) < 2:
+        return None, "At least 2 subject images required"
+    
     payload = {
-        "model_name": "kling-v2.1",  # Best for face preservation
+        "model_name": "kling-v2",
         "prompt": prompt,
-        "image": image_base64,
-        "image_reference": "face",    # Critical for face control
-        "image_fidelity": 0.97,       # 97% similarity
-        "human_fidelity": 0.97,       # 97% facial features
-        "aspect_ratio": "1:1",
-        "n": 1
+        "subject_image_list": subject_image_list,
+        "n": 1,
+        "aspect_ratio": "1:1"
     }
     
     try:
         response = requests.post(CREATE_TASK_ENDPOINT, json=payload, headers=headers)
         response.raise_for_status()
-        return response.json()
-    except Exception as e:
-        logger.error(f"API Error: {str(e)}")
-        return None
+        return response.json(), None
+    except requests.exceptions.RequestException as e:
+        logger.error(f"API request failed: {str(e)}")
+        if hasattr(e, 'response') and e.response:
+            logger.error(f"API response: {e.response.text}")
+        return None, f"API Error: {str(e)}"
 
 def check_task_status(task_id):
+    """Check task completion status"""
     headers = {"Authorization": f"Bearer {generate_jwt_token()}"}
+    status_url = f"{API_BASE_URL}/v1/images/multi-image2image/{task_id}"
+    
     try:
-        response = requests.get(
-            f"{API_BASE_URL}/v1/images/generations/{task_id}",
-            headers=headers
-        )
+        response = requests.get(status_url, headers=headers)
         response.raise_for_status()
-        return response.json()
-    except Exception as e:
-        logger.error(f"Status Check Error: {str(e)}")
-        return None
+        return response.json(), None
+    except requests.exceptions.RequestException as e:
+        return None, f"Status check failed: {str(e)}"
 
-# ===== MAIN FUNCTION =====
-def transform_face(image_path, prompt):
-    """Full transformation workflow"""
-    # Validate image
-    is_valid, error_msg = validate_face_image(image_path)
-    if not is_valid:
-        return None, error_msg
+# ===== MAIN PROCESSING =====
+def generate_image(subject_images, prompt):
+    """Handle complete image generation workflow"""
+    # Validate images
+    for img in subject_images:
+        if img:  # Only validate non-empty images
+            is_valid, error_msg = validate_image(img)
+            if not is_valid:
+                return None, error_msg
     
-    try:
-        # Prepare image
-        image_base64 = prepare_image_base64(image_path)
-        if not image_base64:
-            return None, "Failed to process image"
-        
-        # Create task
-        task_data = create_face_task(image_base64, prompt)
-        if not task_data or task_data.get("code") != 0:
-            return None, "Failed to start transformation"
-        
-        task_id = task_data["data"]["task_id"]
-        logger.info(f"Task created: {task_id}")
+    # Create task
+    task_response, error = create_multi_image_task(subject_images, 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"Task created: {task_id}")
+    
+    # Poll for results (max 10 minutes)
+    for _ in range(60):
+        task_data, error = check_task_status(task_id)
+        if error:
+            return None, error
+            
+        status = task_data["data"]["task_status"]
         
-        # Check results (max 3 minutes)
-        for _ in range(18):  # 18 attempts × 10 seconds
-            time.sleep(10)
-            status_data = check_task_status(task_id)
-            if not status_data:
-                continue
-                
-            if status_data["data"]["task_status"] == "succeed":
-                image_url = status_data["data"]["task_result"]["images"][0]["url"]
-                img_data = requests.get(image_url).content
-                output_path = f"/tmp/face_result_{task_id}.png"
+        if status == "succeed":
+            image_url = task_data["data"]["task_result"]["images"][0]["url"]
+            try:
+                response = requests.get(image_url)
+                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)
-                return output_path, None
+                    f.write(response.content)
+                return str(output_path), None
+            except Exception as e:
+                return None, f"Failed to download result: {str(e)}"
                 
-            elif status_data["data"]["task_status"] in ("failed", "canceled"):
-                error_msg = status_data["data"].get("task_status_msg", "Task failed")
-                return None, error_msg
-        
-        return None, "Processing timed out"
-        
-    except Exception as e:
-        return None, f"Error: {str(e)}"
+        elif status in ("failed", "canceled"):
+            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 =====
-with gr.Blocks(title="Face Transformer") as app:
-    gr.Markdown("# 🎭 Exact Face Transformation (97% Match)")
-    gr.Markdown("Upload ONE face photo for style transformation (97% similarity)")
+def process_interface(subject_image1, subject_image2, subject_image3, subject_image4, prompt):
+    # Filter out None/empty images
+    subject_images = [img for img in [subject_image1, subject_image2, subject_image3, subject_image4] if img]
+    
+    if len(subject_images) < 2:
+        return None, None, "Please upload at least 2 subject images"
+    
+    output_path, error = generate_image(subject_images, prompt)
+    if error:
+        return None, None, error
+    
+    return output_path, output_path, "Generation successful!"
+
+with gr.Blocks(title="Kling AI Multi-Image Generator") as app:
+    gr.Markdown("## 🖼️ Kling AI Multi-Image to Image")
+    gr.Markdown("Combine features from multiple images into one result")
     
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(
-                type="filepath",
-                label="Upload Face Photo",
-                sources=["upload"],
-                height=300
-            )
+            gr.Markdown("### Input Settings")
+            with gr.Row():
+                subject_image1 = gr.Image(type="filepath", label="Subject Image 1 *")
+                subject_image2 = gr.Image(type="filepath", label="Subject Image 2 *")
+            with gr.Row():
+                subject_image3 = gr.Image(type="filepath", label="Subject Image 3 (Optional)")
+                subject_image4 = gr.Image(type="filepath", label="Subject Image 4 (Optional)")
+            
             prompt_input = gr.Textbox(
-                label="Style Prompt", 
-                placeholder="e.g. 'anime character', 'watercolor portrait'"
+                label="Transformation Prompt",
+                placeholder="Describe how to combine these images"
             )
-            generate_btn = gr.Button("Transform", variant="primary")
             
-            gr.Markdown("### Requirements")
+            generate_btn = gr.Button("Generate", variant="primary")
+            
+            gr.Markdown("### Requirements (* = required)")
             gr.Markdown("""
-            - **Single clear face photo**
-            - Front-facing works best
-            - No glasses/masks
-            - Max 10MB (JPG/PNG)
-            - Min 300x300px
+            - **At least 2 subject images** (marked with *)
+            - Max 4 images total
+            - Max size: 10MB per image
+            - Formats: JPG, PNG
+            - Min dimensions: 300x300px
             """)
             
         with gr.Column():
-            output_image = gr.Image(label="Result", interactive=False, height=400)
+            gr.Markdown("### Output")
+            output_image = gr.Image(label="Generated Image", interactive=False, height=400)
             output_file = gr.File(label="Download Result")
             status_output = gr.Textbox(label="Status", interactive=False)
     
+    # Modified inputs to accept individual components
     generate_btn.click(
-        fn=lambda img, prompt: transform_face(img, prompt) + (None,),
-        inputs=[image_input, prompt_input],
+        fn=process_interface,
+        inputs=[subject_image1, subject_image2, subject_image3, subject_image4, prompt_input],
         outputs=[output_image, output_file, status_output]
     )
 
 if __name__ == "__main__":
-    app.launch(server_name="0.0.0.0", server_port=7860)
+    app.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )