diff --git a/__init__.py b/__init__.py index 4031cd4..3e9a7a9 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,5 @@ from .nodes.heygem import HeyGemF2F, HeyGemF2FFromFile -from .nodes.s3 import S3Download, S3Upload +from .nodes.s3 import S3Download, S3Upload, S3UploadURL from .nodes.text import * from .nodes.traverse_folder import TraverseFolder from .nodes.unload_all_models import UnloadAllModels @@ -20,6 +20,7 @@ NODE_CLASS_MAPPINGS = { "COSUpload": COSUpload, "COSDownload": COSDownload, "S3Upload": S3Upload, + "S3UploadURL": S3UploadURL, "S3Download": S3Download, "VideoCutCustom": VideoCut, "VideoCutByFramePoint": VideoCutByFramePoint, @@ -43,6 +44,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "COSUpload": "COS上传", "COSDownload": "COS下载", "S3Upload": "S3上传", + "S3UploadURL": "S3上传-返回URL", "S3Download": "S3下载", "VideoCutCustom": "视频剪裁", "VideoCutByFramePoint": "视频剪裁(精确帧位)", diff --git a/nodes/s3.py b/nodes/s3.py index ecf5a4d..1c82375 100644 --- a/nodes/s3.py +++ b/nodes/s3.py @@ -95,5 +95,52 @@ class S3Upload: return (dest_key,) +class S3UploadURL: + """AWS S3上传 返回URL""" + + @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 = "不忘科技-自定义节点🚩/S3" + + def upload(self, s3_bucket, path, subfolder): + loguru.logger.info(f"S3 UPLOAD {path} to {s3_bucket}/{subfolder}") + try: + if "aws_key_id" in list(os.environ.keys()): + yaml_config = os.environ + else: + 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}") + url = f"https://{s3_bucket}.s3.dualstack.ap-northeast-2.amazonaws.com/{dest_key}" + return (url,) + +