From 52f23c505c7f5d3bf710996424f97dae605027d2 Mon Sep 17 00:00:00 2001 From: binary-husky Date: Sun, 17 Nov 2024 17:45:53 +0000 Subject: [PATCH] media-gpt update --- config.py | 2 +- crazy_functions/VideoResource_GPT.py | 4 +- shared_utils/docker_as_service_api.py | 70 +++++++++++++++++++++++++++ 3 files changed, 73 insertions(+), 3 deletions(-) create mode 100644 shared_utils/docker_as_service_api.py diff --git a/config.py b/config.py index e71a59b1..2b64a9ff 100644 --- a/config.py +++ b/config.py @@ -308,7 +308,7 @@ PLUGIN_HOT_RELOAD = False # 自定义按钮的最大数量限制 NUM_CUSTOM_BASIC_BTN = 4 - +DAAS_SERVER_URL = "http://localhost:48000/stream" """ diff --git a/crazy_functions/VideoResource_GPT.py b/crazy_functions/VideoResource_GPT.py index 3272b6d3..156184b0 100644 --- a/crazy_functions/VideoResource_GPT.py +++ b/crazy_functions/VideoResource_GPT.py @@ -64,7 +64,7 @@ def download_video(bvid, user_name, chatbot, history): # download audio chatbot.append((None, "下载音频, 请稍等...")); yield from update_ui(chatbot=chatbot, history=history) - downloaded_files = download_bilibili(bvid, only_audio=True, user_name=user_name) + downloaded_files = yield from download_bilibili(bvid, only_audio=True, user_name=user_name, chatbot=chatbot, history=history) # preview preview_list = [promote_file_to_downloadzone(fp) for fp in downloaded_files] @@ -81,7 +81,7 @@ def download_video(bvid, user_name, chatbot, history): # download video chatbot.append((None, "下载视频, 请稍等...")); yield from update_ui(chatbot=chatbot, history=history) - downloaded_files_part2 = download_bilibili(bvid, only_audio=False, user_name=user_name) + downloaded_files_part2 = yield from download_bilibili(bvid, only_audio=False, user_name=user_name, chatbot=chatbot, history=history) # preview preview_list = [promote_file_to_downloadzone(fp) for fp in downloaded_files_part2] diff --git a/shared_utils/docker_as_service_api.py b/shared_utils/docker_as_service_api.py new file mode 100644 index 00000000..896085ee --- /dev/null +++ b/shared_utils/docker_as_service_api.py @@ -0,0 +1,70 @@ +import requests +import pickle +import io +import os +from pydantic import BaseModel, Field +from typing import Optional, Dict +from loguru import logger + +class DockerServiceApiComModel(BaseModel): + client_command: Optional[str] = Field(default=None, title="Client command", description="The command to be executed on the client side") + client_file_attach: Optional[dict] = Field(default=None, title="Client file attach", description="The file to be attached to the client side") + server_message: Optional[str] = Field(default=None, title="Server standard error", description="The standard error from the server side") + server_std_err: Optional[str] = Field(default=None, title="Server standard error", description="The standard error from the server side") + server_std_out: Optional[str] = Field(default=None, title="Server standard output", description="The standard output from the server side") + server_file_attach: Optional[dict] = Field(default=None, title="Server file attach", description="The file to be attached to the server side") + +def process_received(received: DockerServiceApiComModel, save_file_dir="./daas_output", output_manifest=None): + # Process the received data + if received.server_message: + output_manifest['server_message'] += received.server_message + if received.server_std_err: + output_manifest['server_std_err'] += received.server_std_err + if received.server_std_out: + output_manifest['server_std_out'] += received.server_std_out + if received.server_file_attach: + # print(f"Recv file attach: {received.server_file_attach}") + for file_name, file_content in received.server_file_attach.items(): + new_fp = os.path.join(save_file_dir, file_name) + new_fp_dir = os.path.dirname(new_fp) + if not os.path.exists(new_fp_dir): + os.makedirs(new_fp_dir, exist_ok=True) + with open(new_fp, 'wb') as f: + f.write(file_content) + output_manifest['server_file_attach'].append(new_fp) + return output_manifest + +def stream_daas(docker_service_api_com_model, server_url, save_file_dir): + # Prepare the file + # Pickle the object + pickled_data = pickle.dumps(docker_service_api_com_model) + + # Create a file-like object from the pickled data + file_obj = io.BytesIO(pickled_data) + + # Prepare the file for sending + files = {'file': ('docker_service_api_com_model.pkl', file_obj, 'application/octet-stream')} + + # Send the POST request + response = requests.post(server_url, files=files, stream=True) + + max_full_package_size = 1024 * 1024 * 1024 * 1 # 1 GB + + received_output_manifest = {} + received_output_manifest['server_message'] = "" + received_output_manifest['server_std_err'] = "" + received_output_manifest['server_std_out'] = "" + received_output_manifest['server_file_attach'] = [] + + # Check if the request was successful + if response.status_code == 200: + # Process the streaming response + for chunk in response.iter_content(max_full_package_size): + if chunk: + received = pickle.loads(chunk) + received_output_manifest = process_received(received, save_file_dir, output_manifest = received_output_manifest) + yield received_output_manifest + else: + logger.error(f"Error: Received status code {response.status_code}, response.text: {response.text}") + + return received_output_manifest \ No newline at end of file