ADD 增加S3上传返回URL接口

__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": "视频剪裁(精确帧位)",

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,)
+
+