ADD 增加上传图片返回url接口

nodes/object_storage_nodes.py

@@ -2,9 +2,12 @@ import os
 
 import boto3
 import loguru
+import torch
 import yaml
 from qcloud_cos import CosConfig, CosS3Client, CosClientError, CosServiceError
 
+from ..utils.image_utils import tensor_to_tempfile
+
 
 class COSDownload:
     """腾讯云COS下载"""
@@ -235,7 +238,7 @@ class S3UploadURL:
         }
 
     RETURN_TYPES = ("STRING",)
-    RETURN_NAMES = ("S3文件Key",)
+    RETURN_NAMES = ("URL",)
 
     FUNCTION = "upload"
     CATEGORY = "不忘科技-自定义节点🚩/对象存储/S3"
@@ -267,3 +270,51 @@ class S3UploadURL:
             raise Exception(f"S3上传失败！ bucket {s3_bucket}; local_path {path}; subfolder {subfolder}")
         url = f"https://cdn.roasmax.cn/{dest_key}"
         return (url,)
+
+
+class S3UploadIMAGEURL:
+    """AWS S3上传 返回URL"""
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", {"multiline": True}),
+                "subfolder": ("STRING", {"default": "test"}),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("URL",)
+
+    FUNCTION = "upload"
+    CATEGORY = "不忘科技-自定义节点🚩/对象存储/S3"
+
+    def upload(self, image:torch.Tensor, subfolder):
+        s3_bucket = "modal-media-cache"
+        loguru.logger.info(f"S3 UPLOAD image to {s3_bucket}/{subfolder}")
+        path = tensor_to_tempfile(image).name
+        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://cdn.roasmax.cn/{dest_key}"
+        return (url,)