ADD 增加S3上传下载节点
This commit is contained in:
parent
7babcd2869
commit
f506eb538b
|
|
@ -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下载",
|
||||
|
|
|
|||
85
nodes/cos.py
85
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),
|
||||
)))
|
||||
if os.sep in cos_key or "/" in cos_key or "\\" in cos_key:
|
||||
os.makedirs(
|
||||
os.path.join(
|
||||
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||||
"download",
|
||||
os.path.dirname(cos_key),
|
||||
),
|
||||
exist_ok=True,
|
||||
)
|
||||
os.path.dirname(cos_key_in),
|
||||
os.path.basename(cos_key_in)
|
||||
)
|
||||
loguru.logger.info(f"COS DOWNLOAD to {destination}")
|
||||
os.makedirs(
|
||||
os.path.dirname(destination),
|
||||
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(
|
||||
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||||
"download",
|
||||
os.path.dirname(cos_key),
|
||||
os.path.basename(cos_key),
|
||||
),
|
||||
DestFilePath=destination
|
||||
)
|
||||
break
|
||||
except CosClientError or CosServiceError as e:
|
||||
print(f"下载失败 {e}")
|
||||
raise Exception(f"COS下载失败！ bucket {cos_bucket}; key {cos_key}")
|
||||
return (
|
||||
os.path.join(
|
||||
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||||
"download",
|
||||
os.path.dirname(cos_key),
|
||||
os.path.basename(cos_key),
|
||||
),
|
||||
destination,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -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(
|
||||
[
|
||||
yaml_config["subfolder"],
|
||||
(
|
||||
path.split("/")[-1]
|
||||
if "/" in path
|
||||
else path.split("\\")[-1]
|
||||
),
|
||||
]
|
||||
),
|
||||
Key=dest_key,
|
||||
LocalFilePath=path,
|
||||
)
|
||||
break
|
||||
except CosClientError or CosServiceError as e:
|
||||
raise RuntimeError("上传失败")
|
||||
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状态失败")
|
||||
raise Exception(f"COS上传失败 bucket {cos_bucket}; local_path {path}; subfolder {subfolder}")
|
||||
return (
|
||||
"/".join(
|
||||
[
|
||||
yaml_config["subfolder"],
|
||||
path.split("/")[-1] if "/" in path else path.split("\\")[-1],
|
||||
]
|
||||
),
|
||||
dest_key,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,)
|
||||
|
||||
|
||||
|
||||
|
||||
|
|
@ -7,3 +7,4 @@ ultralytics
|
|||
cos-python-sdk-v5
|
||||
sqlalchemy
|
||||
tencentcloud-sdk-python
|
||||
boto3
|
||||
Loading…
Reference in New Issue