bowong
/
ComfyUI-CustomNode
6
0
Fork 0
Code Pull Requests Releases Activity

PERF cos上传下载可选bucket

This commit is contained in:
kyj@bowong.ai 2025-03-13 17:40:31 +08:00
parent be21b58126
commit cc0b3da745
1 changed files with 16 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
nodes/cos.py
View File

@ -2,6 +2,7 @@ import json
import os import os
import urllib import urllib
import loguru
import yaml import yaml
from qcloud_cos import CosConfig, CosS3Client, CosClientError, CosServiceError from qcloud_cos import CosConfig, CosS3Client, CosClientError, CosServiceError
@ -13,6 +14,7 @@ class COSDownload:
def INPUT_TYPES(s): def INPUT_TYPES(s):
return { return {
"required": { "required": {
"cos_bucket": ("STRING", {"default": "bwkj-cos-1324682537"}),
"cos_key": ("STRING", {"multiline": True}), "cos_key": ("STRING", {"multiline": True}),
} }
} }
@ -22,11 +24,16 @@ class COSDownload:
FUNCTION = "download" FUNCTION = "download"
CATEGORY = "不忘科技-自定义节点🚩" CATEGORY = "不忘科技-自定义节点🚩"
def download(self, cos_key): def download(self, cos_bucket, cos_key):
loguru.logger.info("download to {}".format(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: if os.sep in cos_key or "/" in cos_key or "\\" in cos_key:
os.makedirs( os.makedirs(
os.path.join( os.path.join(
os.path.dirname(os.path.abspath(__file__)), os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
"download", "download",
os.path.dirname(cos_key), os.path.dirname(cos_key),
), ),
@ -49,10 +56,10 @@ class COSDownload:
) )
client = CosS3Client(config) client = CosS3Client(config)
response = client.download_file( response = client.download_file(
Bucket=yaml_config["bucket"], Bucket=cos_bucket,
Key=cos_key, Key=cos_key,
DestFilePath=os.path.join( DestFilePath=os.path.join(
os.path.dirname(os.path.abspath(__file__)), os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
"download", "download",
os.path.dirname(cos_key), os.path.dirname(cos_key),
os.path.basename(cos_key), os.path.basename(cos_key),
@ -63,7 +70,7 @@ class COSDownload:
print(f"下载失败 {e}") print(f"下载失败 {e}")
return ( return (
os.path.join( os.path.join(
os.path.dirname(os.path.abspath(__file__)), os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
"download", "download",
os.path.dirname(cos_key), os.path.dirname(cos_key),
os.path.basename(cos_key), os.path.basename(cos_key),
@ -78,6 +85,7 @@ class COSUpload:
def INPUT_TYPES(s): def INPUT_TYPES(s):
return { return {
"required": { "required": {
"cos_bucket": ("STRING", {"default": "bwkj-cos-1324682537"}),
"path": ("STRING", {"multiline": True}), "path": ("STRING", {"multiline": True}),
} }
} }
@ -88,7 +96,7 @@ class COSUpload:
FUNCTION = "upload" FUNCTION = "upload"
CATEGORY = "不忘科技-自定义节点🚩" CATEGORY = "不忘科技-自定义节点🚩"
def upload(self, path): def upload(self, cos_bucket, path):
for i in range(0, 10): for i in range(0, 10):
try: try:
with open( with open(
@ -106,7 +114,7 @@ class COSUpload:
) )
client = CosS3Client(config) client = CosS3Client(config)
response = client.upload_file( response = client.upload_file(
Bucket=yaml_config["bucket"], Bucket=cos_bucket,
Key="/".join( Key="/".join(
[ [
yaml_config["subfolder"], yaml_config["subfolder"],
@ -123,7 +131,7 @@ class COSUpload:
except CosClientError or CosServiceError as e: except CosClientError or CosServiceError as e:
raise RuntimeError("上传失败") raise RuntimeError("上传失败")
data = {"prompt_id": "", data = {"prompt_id": "",
"video_url": "https://{}.cos.{}.myqcloud.com/{}".format(yaml_config['bucket'], yaml_config['region'], "video_url": "https://{}.cos.{}.myqcloud.com/{}".format(cos_bucket, yaml_config['region'],
'/'.join([yaml_config['subfolder'], '/'.join([yaml_config['subfolder'],
path.split('/')[ path.split('/')[
-1] if '/' in path else -1] if '/' in path else