From 8add142be0bb20ad34fba5019e8d118c7880bce1 Mon Sep 17 00:00:00 2001 From: MaraScott Date: Fri, 15 Nov 2024 17:31:01 +0100 Subject: [PATCH] add Mara Load Image --- MaraScott_Nodes.py | 7 +- py/nodes/Image/LoadImage_v1.py | 122 +++++++++++++++++++++++++++++++++ py/utils/helper.py | 6 +- 3 files changed, 133 insertions(+), 2 deletions(-) create mode 100644 py/nodes/Image/LoadImage_v1.py diff --git a/MaraScott_Nodes.py b/MaraScott_Nodes.py index 2a2cb45..46612f9 100644 --- a/MaraScott_Nodes.py +++ b/MaraScott_Nodes.py @@ -7,6 +7,7 @@ # ### +from .py.nodes.Image.LoadImage_v1 import LoadImage_v1 from .py.nodes.Bus.AnyBus_v2 import AnyBus_v2 from .py.nodes.Info.DisplayInfo_v2 import DisplayInfo_v2 from .py.nodes.UpscalerRefiner.McBoaty_v3 import UpscalerRefiner_McBoaty_v3 @@ -61,6 +62,8 @@ "MaraScottGetModelBlocks_v1": GetModelBlocks_v1, + "MaraScottLoadImage_v1": LoadImage_v1, + "MaraScott_Kijai_TokenCounter_v1": TokenCounter_v1, "MaraScott_laksjdjf_Hires_v1": Hires_v1, } @@ -77,7 +80,7 @@ "MaraScottMcBoatyUpscaler_v4": "\ud83d\udc30 Upscaler - McBoaty [1/3] v4 /u", "MaraScottMcBoatyTilePrompter_v4": "\ud83d\udc30 Tile Prompter - McBoaty [2/3] v4 /u", "MaraScottMcBoatyRefiner_v4": "\ud83d\udc30 Refiner - McBoaty [3/3] v4 /u", - "MaraScottMcBoatyUpscalerRefinerNode_v3": "\ud83d\udc30 Large Refiner - McBoaty v3 /u", + "MaraScottUpscalerRefinerNode_v3": "\ud83d\udc30 Large Refiner - McBoaty v3 /u", "MaraScottSetInpaintingByMask_v1": "\ud83d\udc30 Set Inpainting Tile by mask - McInpainty [1/2] v1 /m", "MaraScottPasteInpaintingByMask_v1": "\ud83d\udc30 Paste Inpainting Tile by mask - McInpainty [2/2] v1 /m", @@ -101,6 +104,8 @@ "MaraScottGetModelBlocks_v1": "\ud83d\udc30 Get Model Blocks v1 /b", + "MaraScottLoadImage_v1": "\ud83d\udc30 Load Image v1 /i", + "MaraScott_Kijai_TokenCounter_v1": "\ud83d\udc30 TokenCounter (from kijai/ComfyUI-KJNodes) /v", "MaraScott_laksjdjf_Hires_v1": "\ud83d\udc30 Apply Kohya's HiresFix (from laksjdjf) sd1.5 only /sd15", diff --git a/py/nodes/Image/LoadImage_v1.py b/py/nodes/Image/LoadImage_v1.py new file mode 100644 index 0000000..d74cf86 --- /dev/null +++ b/py/nodes/Image/LoadImage_v1.py @@ -0,0 +1,122 @@ +import os +import hashlib +import nodes +import folder_paths + +from pathlib import Path + +from ...utils.constants import get_name, get_category +from ...utils.log import log +from ...utils.helper import current_method + +from inspect import currentframe as cf + + +class nodes_LoadImage_v1(nodes.LoadImage): + + def __init__(self): + super().__init__() + + ALLOWED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tiff", ".webp"} + + @classmethod + def INPUT_TYPES(s): + + input_dir = folder_paths.get_input_directory() + exclude_folders = [] + file_list = [] + for root, dirs, files in os.walk(input_dir): + # Exclude specific folders + dirs[:] = [d for d in dirs if d not in exclude_folders] + + for file in files: + if os.path.splitext(file)[1].lower() in s.ALLOWED_EXTENSIONS: + relpath = os.path.relpath(os.path.join(root, file), start=input_dir) + # fix for windows + relpath = relpath.replace("\\", "/") + file_list.append(relpath) + + return { + "required": { + "image": (sorted(file_list), {"image_upload": True}) + }, + } + + @classmethod + def _INPUT_TYPES(s, EXTRA_INPUT_TYPES): + _input_types = getattr(super(), "INPUT_TYPES", lambda: {})() + + for key, value in EXTRA_INPUT_TYPES.items(): + if key in _input_types: + if isinstance(_input_types[key], dict) and isinstance(value, dict): + _input_types[key].update(value) + else: + _input_types[key] = value + else: + _input_types[key] = value + + return _input_types + + @classmethod + def _RETURN_OUTPUTS(s, OUTPUT_PARAM_NAME): + OUTPUT_PARAMS = getattr(super(), OUTPUT_PARAM_NAME, ()) + if OUTPUT_PARAM_NAME == "RETURN_NAMES" and len(OUTPUT_PARAMS) == 0: + OUTPUT_PARAMS = getattr(super(), "RETURN_TYPES", ()) + if not isinstance(OUTPUT_PARAMS, tuple): + OUTPUT_PARAMS = (OUTPUT_PARAMS,) + return OUTPUT_PARAMS + + @classmethod + def IS_CHANGED(s, **kwargs): + image = kwargs.get("image", None) + return current_method(super(), cf())(image) + + @classmethod + def VALIDATE_INPUTS(s, **kwargs): + image = kwargs.get("image", None) + return current_method(super(), cf())(image) + + +class LoadImage_v1(nodes_LoadImage_v1): + + def __init__(self): + super().__init__() + + NAME = get_name("Load Image") + + @classmethod + def _INPUT_TYPES(s): + + INPUT_TYPES = { + "required": { + "filename_with_ext": ("BOOLEAN", {"default": False}), + }, + "hidden": {}, + } + + return super()._INPUT_TYPES(INPUT_TYPES) + + RETURN_TYPES = nodes_LoadImage_v1()._RETURN_OUTPUTS("RETURN_TYPES") + ("STRING",) + RETURN_NAMES = nodes_LoadImage_v1()._RETURN_OUTPUTS("RETURN_NAMES") + ("File Name",) + + INPUT_IS_LIST = False + FUNCTION = "fn" + CATEGORY = get_category("image") + OUTPUT_NODE = True + OUTPUT_IS_LIST = (None,) * len(RETURN_TYPES) + + def fn(s, **kwargs): + image = kwargs.get("image", None) + filename_with_ext = kwargs.get("filename_with_ext", None) + + file_path = folder_paths.get_annotated_filepath(image) + output_image, output_mask = super().load_image(file_path) + filename = file_path.replace(folder_paths.get_input_directory() + "\\", "") + if not filename_with_ext: + filename = os.path.splitext(filename)[0] + + return ( + output_image, + output_mask, + filename, + ) diff --git a/py/utils/helper.py b/py/utils/helper.py index d3a5e4b..a447343 100644 --- a/py/utils/helper.py +++ b/py/utils/helper.py @@ -48,4 +48,8 @@ def analyze_object(obj): return {"properties": properties, "methods": methods} def natural_key(entry): - return [int(text) if text.isdigit() else text for text in re.split(r'(\d+)', entry)] \ No newline at end of file + return [int(text) if text.isdigit() else text for text in re.split(r'(\d+)', entry)] + +def current_method(super_class, current_frame): + return getattr(super_class, current_frame.f_code.co_name) + \ No newline at end of file