Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import requests | |
| import base64 | |
| from pathlib import Path | |
| import jwt | |
| import time | |
| import os | |
| # === Kling AI API configuration === | |
| ACCESS_KEY_ID = "AGBGmadNd9hakFYfahytyQQJtN8CJmDJ" | |
| ACCESS_KEY_SECRET = "dp3pAe4PpdmnAHCAPgEd3PyLmBQrkMde" | |
| API_URL = "https://api-singapore.klingai.com/v1/generate/image" | |
| # === Generate JWT Token === | |
| def generate_jwt_token(): | |
| headers = { | |
| "alg": "HS256", | |
| "typ": "JWT" | |
| } | |
| now = int(time.time()) | |
| payload = { | |
| "iss": ACCESS_KEY_ID, | |
| "exp": now + 1800, # valid for 30 mins | |
| "nbf": now - 5 | |
| } | |
| return jwt.encode(payload, ACCESS_KEY_SECRET, algorithm="HS256", headers=headers) | |
| # === Image to Image Generation === | |
| def generate_image(reference_image, prompt=""): | |
| if not reference_image: | |
| return None, None | |
| try: | |
| with open(reference_image, "rb") as img_file: | |
| reference_base64 = base64.b64encode(img_file.read()).decode("utf-8") | |
| except Exception as e: | |
| return None, f"Error reading reference image: {str(e)}" | |
| headers = { | |
| "Authorization": f"Bearer {generate_jwt_token()}", | |
| "Content-Type": "application/json" | |
| } | |
| payload = { | |
| "reference_image": reference_base64, | |
| "mode": "face", # face reference mode | |
| "prompt": prompt or "Match face with high fidelity", | |
| "strength": 0.97, | |
| "output_format": "png" | |
| } | |
| try: | |
| # Start generation | |
| response = requests.post(API_URL, json=payload, headers=headers, timeout=30) | |
| response.raise_for_status() | |
| data = response.json() | |
| task_id = data.get("task_id") or data.get("id") | |
| if not task_id: | |
| return None, f"Error: No task ID returned: {data}" | |
| # Poll for completion | |
| status_url = f"https://api-singapore.klingai.com/v1/predictions/{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() | |
| status = status_data.get("status") | |
| if status == "succeeded": | |
| # Find image URL in all possible keys | |
| image_url = ( | |
| status_data.get("image_url") | |
| or status_data.get("result", {}).get("image_url") | |
| or (status_data.get("output")[0] if isinstance(status_data.get("output"), list) else status_data.get("output")) | |
| ) | |
| if not image_url: | |
| return None, f"Error: No image URL found in response: {status_data}" | |
| # Download the generated image | |
| img_resp = requests.get(image_url, timeout=30) | |
| img_resp.raise_for_status() | |
| output_path = Path("output.png") | |
| with open(output_path, "wb") as f: | |
| f.write(img_resp.content) | |
| return str(output_path), None | |
| elif status == "failed": | |
| return None, "Error: Generation failed." | |
| time.sleep(5) | |
| return None, "Error: Timed out." | |
| except requests.exceptions.RequestException as e: | |
| return None, f"Request error: {str(e)}" | |
| # === Gradio Interface === | |
| def process(reference_image, prompt): | |
| image_path, error = generate_image(reference_image, prompt) | |
| if error: | |
| return None, None, error | |
| download_link = image_path if image_path else None | |
| return image_path, download_link, "✅ Image generated successfully!" | |
| iface = gr.Interface( | |
| fn=process, | |
| inputs=[ | |
| gr.Image(type="filepath", label="Upload Reference Face Image"), | |
| gr.Textbox(lines=2, placeholder="Optional prompt", label="Prompt") | |
| ], | |
| outputs=[ | |
| gr.Image(label="Generated Image"), | |
| gr.File(label="Download Image"), | |
| gr.Textbox(label="Status") | |
| ], | |
| title="Kling AI - Single Reference Face Generator", | |
| description="Upload a reference face image. The model will generate an image based on that face with 97% strength." | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch(server_name="0.0.0.0", server_port=7860) | |