ADD 增加S3上传下载节点

__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下载",

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),
+                    os.path.dirname(cos_key_in),
-                )))
+                    os.path.basename(cos_key_in)
-        if os.sep in cos_key or "/" in cos_key or "\\" in cos_key:
+                )
-            os.makedirs(
+        loguru.logger.info(f"COS DOWNLOAD to {destination}")
-                os.path.join(
+        os.makedirs(
-                    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+            os.path.dirname(destination),
-                    "download",
+            exist_ok=True,
-                    os.path.dirname(cos_key),
+        )
-                ),
-                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(
+                    DestFilePath=destination
-                        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
-                        "download",
-                        os.path.dirname(cos_key),
-                        os.path.basename(cos_key),
-                    ),
                 )
                 break
             except CosClientError or CosServiceError as e:
-                print(f"下载失败 {e}")
+                raise Exception(f"COS下载失败！ bucket {cos_bucket}; key {cos_key}")
         return (
-            os.path.join(
+            destination,
-                os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
-                "download",
-                os.path.dirname(cos_key),
-                os.path.basename(cos_key),
-            ),
         )
 
 
@@ -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(
+                    Key=dest_key,
-                        [
-                            yaml_config["subfolder"],
-                            (
-                                path.split("/")[-1]
-                                if "/" in path
-                                else path.split("\\")[-1]
-                            ),
-                        ]
-                    ),
                     LocalFilePath=path,
                 )
                 break
             except CosClientError or CosServiceError as e:
-                raise RuntimeError("上传失败")
+                raise Exception(f"COS上传失败 bucket {cos_bucket}; local_path {path}; subfolder {subfolder}")
-        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状态失败")
         return (
-            "/".join(
+            dest_key,
-                [
-                    yaml_config["subfolder"],
-                    path.split("/")[-1] if "/" in path else path.split("\\")[-1],
-                ]
-            ),
         )

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

requirements.txt

@@ -6,4 +6,5 @@ opencv-python
 ultralytics
 cos-python-sdk-v5
 sqlalchemy
-tencentcloud-sdk-python
+tencentcloud-sdk-python
+boto3