From a5159aa09772a92b170f44dbc8a417de6d249c17 Mon Sep 17 00:00:00 2001
From: "kyj@bowong.ai" <kyj@bowong.ai>
Date: Wed, 11 Jun 2025 15:54:06 +0800
Subject: [PATCH] =?UTF-8?q?ADD=20=E5=A2=9E=E5=8A=A0S3=E4=B8=8A=E4=BC=A0?=
 =?UTF-8?q?=E8=BF=94=E5=9B=9EURL=E6=8E=A5=E5=8F=A3?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 __init__.py |  4 +++-
 nodes/s3.py | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

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