ADD 增加视频帧位计算、字符串判空节点

PERF 重构项目结构
2025-03-04 16:50:54 +08:00
parent e1849bb8b9
commit cafa224374
14 changed files with 718 additions and 650 deletions
--- a/nodes/cos.py
+++ b/nodes/cos.py
@@ -0,0 +1,145 @@
+import json
+import os
+import urllib
+
+import yaml
+from qcloud_cos import CosConfig, CosS3Client, CosClientError, CosServiceError
+
+
+class COSDownload:
+    """腾讯云COS下载"""
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "cos_key": ("STRING", {"multiline": True}),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("视频存储路径",)
+    FUNCTION = "download"
+    CATEGORY = "不忘科技-自定义节点🚩"
+
+    def download(self, 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.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=yaml_config["bucket"],
+                    Key=cos_key,
+                    DestFilePath=os.path.join(
+                        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.abspath(__file__)),
+                "download",
+                os.path.dirname(cos_key),
+                os.path.basename(cos_key),
+            ),
+        )
+
+
+class COSUpload:
+    """腾讯云COS上传"""
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "path": ("STRING", {"multiline": True}),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("COS文件Key",)
+
+    FUNCTION = "upload"
+    CATEGORY = "不忘科技-自定义节点🚩"
+
+    def upload(self, 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=yaml_config["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(yaml_config['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],
+                ]
+            ),
+        )