feat: add kling (#331)

* feat: add kling

Author: timerring

bilive.toml

@@ -40,6 +40,7 @@ sensenova_api_key = "" # Apply for your own SenseNova API key at https://console
 generate_cover = false # whether to generate cover
 cover_prompt = "" # Write your own cover prompt here
 image_gen_model = "minimax" # the image generation model, can be "minimax" or "siliconflow" or "tencent" or "baidu" or "stability" or "luma" or "ideogram" or "recraft" or "amazon" or "hidream"
+dmx_api_token = "" # Apply for your own dmxapi token at https://www.dmxapi.cn/token
 minimax_api_key = "" # Apply for your own Minimax API key at https://platform.minimaxi.com/user-center/basic-information/interface-key
 siliconflow_api_key = "" # Apply for your own SiliconFlow API key at https://cloud.siliconflow.cn/i/3Szr5BVg
 tencent_secret_id = "" # Apply for your own Tencent Cloud API key at https://console.cloud.tencent.com/cam/capi

src/config.py

@@ -92,6 +92,6 @@ def get_interface_config():
 AWS_ACCESS_KEY_ID = config.get("cover", {}).get("aws_access_key_id")
 AWS_SECRET_ACCESS_KEY = config.get("cover", {}).get("aws_secret_access_key")
 HIDREAM_API_KEY = config.get("cover", {}).get("hidream_api_key")
-
+DMX_API_TOKEN = config.get("cover", {}).get("dmx_api_token")
 SLICE_PROMPT = config.get("slice", {}).get("slice_prompt")
 COVER_PROMPT = config.get("cover", {}).get("cover_prompt")

src/cover/image_model_sdk/kling_sdk.py

@@ -0,0 +1,77 @@
+import http.client
+import json
+import base64
+import requests
+import time
+import os
+from src.config import DMX_API_TOKEN, COVER_PROMPT
+
+
+conn = http.client.HTTPSConnection("www.dmxapi.cn")
+
+def get_image_base64(image_path):
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode("utf-8")
+
+def generate_cover(your_file_path):
+    image_base = get_image_base64(your_file_path)
+    payload = json.dumps({
+    "model_name": "kling-v1-5",
+    "prompt": COVER_PROMPT,
+    "image": image_base,
+    "image_reference": "subject"
+    })
+    headers = {
+    'Authorization': f'Bearer {DMX_API_TOKEN}',
+    'Content-Type': 'application/json'
+    }
+    conn.request("POST", "/kling/v1/images/generations?=null", payload, headers)
+    res = conn.getresponse()
+    json_data = json.loads(res.read().decode("utf-8"))
+    print(json_data)
+    return json_data['data']['task_id']
+
+def get_kling_cover(task_id):
+    action = "images"
+    action2 = "generations"
+
+    query_path = f"/kling/v1/{action}/{action2}/{task_id}"
+
+    headers = {
+    'Authorization': f'Bearer {DMX_API_TOKEN}'
+    }
+
+    conn.request("GET", query_path, None, headers)
+    res = conn.getresponse()
+    json_data = json.loads(res.read().decode("utf-8"))
+    if json_data['data']['task_status'] == "succeed":
+        return json_data['data']['task_result']['images'][0]['url']
+    else: 
+        return None
+
+def kling_generate_cover(your_file_path):
+    task_id = generate_cover(your_file_path)
+    start_time = time.time()
+    timeout = 60
+
+    while True:
+        image_url = get_kling_cover(task_id)
+        if image_url is not None:
+            img_data = requests.get(image_url).content
+            cover_name = time.strftime("%Y%m%d%H%M%S") + ".png"
+            temp_cover_path = os.path.join(os.path.dirname(your_file_path), cover_name)
+            with open(temp_cover_path, "wb") as handler:
+                handler.write(img_data)
+            os.remove(your_file_path)
+            return temp_cover_path
+        
+        if time.time() - start_time > timeout:
+            print(f"Generate cover {timeout} seconds timeout")
+            return None
+
+        time.sleep(1)
+        print(f"Waiting for cover generation, {int(time.time() - start_time)} seconds", flush=True)
+
+if __name__ == "__main__":
+    your_file_path = ""
+    print(kling_generate_cover(your_file_path))

tests/test_cover.py

@@ -9,7 +9,7 @@
 from src.cover.image_model_sdk.recraft_sdk import recraft_generate_cover
 from src.cover.image_model_sdk.stability_sdk import stable_diffusion_generate_cover
 from src.cover.image_model_sdk.tencent_sdk import hunyuan_generate_cover
-
+from src.cover.image_model_sdk.kling_sdk import kling_generate_cover
 
 class BaseTest(unittest.TestCase):
     file_path = "your_image_path"
@@ -76,5 +76,11 @@ def test_hunyuan_generate_cover(self):
         self.assertIsNotNone(cover_path)
         print(cover_path, flush=True)
 
+class TestKlingMain(BaseTest):
+    def test_kling_generate_cover(self):
+        cover_path = kling_generate_cover(self.file_path)
+        self.assertIsNotNone(cover_path)
+        print(cover_path, flush=True)
+
 if __name__ == "__main__":
     unittest.main()