diff --git a/__init__.py b/__init__.py index 0b61d5a..17176ce 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,4 @@ +from .nodes.s3 import S3Download, S3Upload from .nodes.text import * from .nodes.traverse_folder import TraverseFolder from .nodes.unload_all_models import UnloadAllModels @@ -17,6 +18,8 @@ NODE_CLASS_MAPPINGS = { "FaceExtract": FaceExtract, "COSUpload": COSUpload, "COSDownload": COSDownload, + "S3Upload": S3Upload, + "S3Download": S3Download, "VideoCutCustom": VideoCut, "VideoCutByFramePoint": VideoCutByFramePoint, "VodToLocal": VodToLocalNode, @@ -35,6 +38,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FaceExtract": "面部提取", "COSUpload": "COS上传", "COSDownload": "COS下载", + "S3Upload": "S3上传", + "S3Download": "S3下载", "VideoCutCustom": "视频剪裁", "VideoCutByFramePoint": "视频剪裁(精确帧位)", "VodToLocal": "腾讯云VOD下载", diff --git a/nodes/cos.py b/nodes/cos.py index a6c3953..755cc94 100644 --- a/nodes/cos.py +++ b/nodes/cos.py @@ -25,20 +25,18 @@ class COSDownload: CATEGORY = "不忘科技-自定义节点🚩" def download(self, cos_bucket, cos_key): - loguru.logger.info("download to {}".format(os.path.join( + cos_key_in = cos_key.replace("/",os.sep) + destination = os.path.join( os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "download", - os.path.dirname(cos_key), - ))) - if os.sep in cos_key or "/" in cos_key or "\\" in cos_key: - os.makedirs( - os.path.join( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - "download", - os.path.dirname(cos_key), - ), - exist_ok=True, - ) + os.path.dirname(cos_key_in), + os.path.basename(cos_key_in) + ) + loguru.logger.info(f"COS DOWNLOAD to {destination}") + os.makedirs( + os.path.dirname(destination), + exist_ok=True, + ) for i in range(0, 10): try: with open( @@ -58,23 +56,13 @@ class COSDownload: response = client.download_file( Bucket=cos_bucket, Key=cos_key, - DestFilePath=os.path.join( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - "download", - os.path.dirname(cos_key), - os.path.basename(cos_key), - ), + DestFilePath=destination ) break except CosClientError or CosServiceError as e: - print(f"下载失败 {e}") + raise Exception(f"COS下载失败！ bucket {cos_bucket}; key {cos_key}") return ( - os.path.join( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - "download", - os.path.dirname(cos_key), - os.path.basename(cos_key), - ), + destination, ) @@ -87,6 +75,7 @@ class COSUpload: "required": { "cos_bucket": ("STRING", {"default": "bwkj-cos-1324682537"}), "path": ("STRING", {"multiline": True}), + "subfolder": ("STRING", {"default": "test"}), } } @@ -96,7 +85,18 @@ class COSUpload: FUNCTION = "upload" CATEGORY = "不忘科技-自定义节点🚩" - def upload(self, cos_bucket, path): + def upload(self, cos_bucket, path, subfolder): + dest_key = "/".join( + [ + subfolder, + ( + path.split("/")[-1] + if "/" in path + else path.split("\\")[-1] + ), + ] + ) + loguru.logger.info(f"COS UPLOAD {path} to {os.path.join(cos_bucket, subfolder)}") for i in range(0, 10): try: with open( @@ -115,39 +115,12 @@ class COSUpload: client = CosS3Client(config) response = client.upload_file( Bucket=cos_bucket, - Key="/".join( - [ - yaml_config["subfolder"], - ( - path.split("/")[-1] - if "/" in path - else path.split("\\")[-1] - ), - ] - ), + Key=dest_key, LocalFilePath=path, ) break except CosClientError or CosServiceError as e: - raise RuntimeError("上传失败") - data = {"prompt_id": "", - "video_url": "https://{}.cos.{}.myqcloud.com/{}".format(cos_bucket, yaml_config['region'], - '/'.join([yaml_config['subfolder'], - path.split('/')[ - -1] if '/' in path else - path.split('\\')[-1], ])) - } - headers = {'Content-Type': 'application/json'} - try: - req = urllib.request.Request("", data=json.dumps(data).encode("utf-8"), headers=headers) - response = urllib.request.urlopen(req) - except: - raise RuntimeError("上报MQ状态失败") + raise Exception(f"COS上传失败 bucket {cos_bucket}; local_path {path}; subfolder {subfolder}") return ( - "/".join( - [ - yaml_config["subfolder"], - path.split("/")[-1] if "/" in path else path.split("\\")[-1], - ] - ), + dest_key, ) diff --git a/nodes/s3.py b/nodes/s3.py new file mode 100644 index 0000000..274ebd9 --- /dev/null +++ b/nodes/s3.py @@ -0,0 +1,92 @@ +import os + +import boto3 +import loguru +import yaml + + +class S3Download: + """AWS S3下载""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "s3_bucket": ("STRING", {"default": "bw-comfyui-input"}), + "s3_key": ("STRING", {"multiline": True}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("视频存储路径",) + FUNCTION = "download" + CATEGORY = "不忘科技-自定义节点🚩" + + def download(self, s3_bucket, s3_key): + s3_key_in = s3_key.replace("/",os.sep) + destination = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "download", + os.path.dirname(s3_key_in), + os.path.basename(s3_key_in), + ) + loguru.logger.info(f"S3 DOWNLOAD to {destination}") + os.makedirs( + os.path.dirname(destination), + exist_ok=True, + ) + try: + with open(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "config.yaml"),encoding="utf-8",mode="r+") as f: + yaml_config = yaml.load(f, Loader=yaml.FullLoader) + client = boto3.client("s3",aws_access_key_id=yaml_config["aws_key_id"],aws_secret_access_key=yaml_config["aws_access_key"]) + client.download_file(s3_bucket, s3_key, destination) + except Exception as e: + raise Exception(f"S3下载失败！ bucket {s3_bucket}; key {s3_key}") + return (destination,) + + +class S3Upload: + """AWS S3上传""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "s3_bucket": ("STRING", {"default": "bw-comfyui-output"}), + "path": ("STRING", {"multiline": True}), + "subfolder": ("STRING", {"default": "test"}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("S3文件Key",) + + FUNCTION = "upload" + CATEGORY = "不忘科技-自定义节点🚩" + + def upload(self, s3_bucket, path, subfolder): + loguru.logger.info(f"S3 UPLOAD {path} to {os.path.join(s3_bucket, subfolder)}") + try: + with open(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "config.yaml"), + encoding="utf-8", mode="r+") as f: + yaml_config = yaml.load(f, Loader=yaml.FullLoader) + client = boto3.client("s3", aws_access_key_id=yaml_config["aws_key_id"], + aws_secret_access_key=yaml_config["aws_access_key"]) + dest_key = "/".join( + [ + subfolder, + ( + path.split("/")[-1] + if "/" in path + else path.split("\\")[-1] + ), + ] + ) + client.upload_file(path, s3_bucket, dest_key) + except Exception as e: + raise Exception(f"S3上传失败！ bucket {s3_bucket}; local_path {path}; subfolder {subfolder}") + return (dest_key,) + + + + diff --git a/requirements.txt b/requirements.txt index 17048c3..a6cfdfb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,5 @@ opencv-python ultralytics cos-python-sdk-v5 sqlalchemy -tencentcloud-sdk-python \ No newline at end of file +tencentcloud-sdk-python +boto3 \ No newline at end of file