diff --git a/__init__.py b/__init__.py index c8abbd8..545fd1c 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,5 @@ +from .nodes.traverse_folder import TraverseFolder +from .nodes.unload_all_models import UnloadAllModels from .nodes.string_empty_judgement import StringEmptyJudgement from .nodes.compute_video_point import VideoStartPointDurationCompute from .nodes.cos import COSUpload, COSDownload @@ -19,6 +21,8 @@ NODE_CLASS_MAPPINGS = { "LogToDB": LogToDB, "VideoPointCompute": VideoStartPointDurationCompute, "StringEmptyJudgement": StringEmptyJudgement, + "unloadAllModels": UnloadAllModels, + "TraverseFolder": TraverseFolder, } # A dictionary that contains the friendly/humanly readable titles for the nodes @@ -32,4 +36,6 @@ NODE_DISPLAY_NAME_MAPPINGS = { "LogToDB": "状态持久化DB", "VideoPointCompute": "视频帧位计算", "StringEmptyJudgement": "字符串是否为空", + "unloadAllModels": "卸载所有已加载模型", + "TraverseFolder": "遍历文件夹" } diff --git a/nodes/compute_video_point.py b/nodes/compute_video_point.py index d4aa267..0cfb560 100644 --- a/nodes/compute_video_point.py +++ b/nodes/compute_video_point.py @@ -11,7 +11,7 @@ def validate_time_format(time_str): def get_duration_wave(audio): waveform, sample_rate = audio["waveform"], audio["sample_rate"] # 防止话说不完 - return ceil(waveform.shape[2] / sample_rate) + 1 + return ceil(waveform.shape[2] / sample_rate) + 0.1 class VideoStartPointDurationCompute: diff --git a/nodes/traverse_folder.py b/nodes/traverse_folder.py new file mode 100644 index 0000000..b2cda48 --- /dev/null +++ b/nodes/traverse_folder.py @@ -0,0 +1,39 @@ +import glob +import os + +class AnyType(str): + """A special class that is always equal in not equal comparisons. Credit to pythongosssss""" + + def __ne__(self, __value: object) -> bool: + return False + + +any = AnyType("*") + +class TraverseFolder: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "folder": ("STRING", {"default": r"E:\comfy\ComfyUI\input\s3", "required": True}), + "subfix": ("STRING", {"default":".mp4", "required": True}), + "recursive": ("BOOLEAN", {"default": True, "required": True}), + "idx": ( + "INT", + {"default": 0, "min": 0, "max": 0xFFFFFF}, + ), + }, + } + + RETURN_TYPES = ("STRING", ) + RETURN_NAMES = ("文件路径",) + + FUNCTION = "compute" + + CATEGORY = "不忘科技-自定义节点🚩" + + def compute(self, folder, subfix, recursive, idx): + files = glob.glob(os.path.join(folder, r"**\*%s" % subfix), recursive=recursive) + if len(files) == 0: + raise RuntimeError("No Files Found") + return (str(files[idx % len(files)]),) \ No newline at end of file diff --git a/nodes/unload_all_models.py b/nodes/unload_all_models.py new file mode 100644 index 0000000..b70ba7c --- /dev/null +++ b/nodes/unload_all_models.py @@ -0,0 +1,32 @@ +import comfy.model_management + +class AnyType(str): + """A special class that is always equal in not equal comparisons. Credit to pythongosssss""" + + def __ne__(self, __value: object) -> bool: + return False + + +any = AnyType("*") + +class UnloadAllModels: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "any":(any,{"forceInput": True}) + }, + "optional": {}, + } + + RETURN_TYPES = () + FUNCTION = "unload_models" + CATEGORY = "不忘科技-自定义节点🚩" + OUTPUT_NODE = True + + def unload_models(self,any=None): + # 卸载所有已加载的模型 + comfy.model_management.soft_empty_cache() + comfy.model_management.unload_all_models() + comfy.model_management.soft_empty_cache() + return () \ No newline at end of file