diff --git a/bilive.toml b/bilive.toml index 136a8eb..f692415 100644 --- a/bilive.toml +++ b/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 diff --git a/src/config.py b/src/config.py index baf81a3..167a42d 100644 --- a/src/config.py +++ b/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") diff --git a/src/cover/image_model_sdk/kling_sdk.py b/src/cover/image_model_sdk/kling_sdk.py new file mode 100644 index 0000000..2dbd61d --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/tests/test_cover.py b/tests/test_cover.py index 527a3eb..08d1132 100644 --- a/tests/test_cover.py +++ b/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()