154 lines
5.5 KiB
Python
154 lines
5.5 KiB
Python
import json
|
|
import os
|
|
import urllib
|
|
|
|
import loguru
|
|
import yaml
|
|
from qcloud_cos import CosConfig, CosS3Client, CosClientError, CosServiceError
|
|
|
|
|
|
class COSDownload:
|
|
"""腾讯云COS下载"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"cos_bucket": ("STRING", {"default": "bwkj-cos-1324682537"}),
|
|
"cos_key": ("STRING", {"multiline": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("视频存储路径",)
|
|
FUNCTION = "download"
|
|
CATEGORY = "不忘科技-自定义节点🚩"
|
|
|
|
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:
|
|
os.makedirs(
|
|
os.path.join(
|
|
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
|
"download",
|
|
os.path.dirname(cos_key),
|
|
),
|
|
exist_ok=True,
|
|
)
|
|
for i in range(0, 10):
|
|
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)
|
|
config = CosConfig(
|
|
Region=yaml_config["region"],
|
|
SecretId=yaml_config["secret_id"],
|
|
SecretKey=yaml_config["secret_key"],
|
|
)
|
|
client = CosS3Client(config)
|
|
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),
|
|
),
|
|
)
|
|
break
|
|
except CosClientError or CosServiceError as e:
|
|
print(f"下载失败 {e}")
|
|
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),
|
|
),
|
|
)
|
|
|
|
|
|
class COSUpload:
|
|
"""腾讯云COS上传"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"cos_bucket": ("STRING", {"default": "bwkj-cos-1324682537"}),
|
|
"path": ("STRING", {"multiline": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("COS文件Key",)
|
|
|
|
FUNCTION = "upload"
|
|
CATEGORY = "不忘科技-自定义节点🚩"
|
|
|
|
def upload(self, cos_bucket, path):
|
|
for i in range(0, 10):
|
|
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)
|
|
config = CosConfig(
|
|
Region=yaml_config["region"],
|
|
SecretId=yaml_config["secret_id"],
|
|
SecretKey=yaml_config["secret_key"],
|
|
)
|
|
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]
|
|
),
|
|
]
|
|
),
|
|
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状态失败")
|
|
return (
|
|
"/".join(
|
|
[
|
|
yaml_config["subfolder"],
|
|
path.split("/")[-1] if "/" in path else path.split("\\")[-1],
|
|
]
|
|
),
|
|
)
|