diff --git a/image_preview.py b/image_preview.py index b72af0f..e081452 100644 --- a/image_preview.py +++ b/image_preview.py @@ -1,13 +1,17 @@ +import hashlib import json import os import random import numpy as np -from PIL import Image +import torch +from PIL import Image, ImageOps, ImageSequence from PIL.PngImagePlugin import PngInfo import folder_paths +import node_helpers from comfy.cli_args import args +from comfy_execution.graph_utils import ExecutionBlocker class JDL_PreviewToLoad: @@ -116,10 +120,100 @@ class JDL_PreviewToLoad: return {"ui": {"images": results, "input_filename": [input_filename]}} +class JDL_LoadImage: + """Load an image from the input directory with an active switch to skip downstream execution.""" + + @classmethod + def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] + files = folder_paths.filter_files_content_types(files, ["image"]) + return { + "required": { + "image": (sorted(files), {"image_upload": True}), + "active": ("BOOLEAN", {"default": True}), + }, + } + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "load_image" + CATEGORY = "utils/image" + + def load_image(self, image, active): + if not active: + return (ExecutionBlocker(None), ExecutionBlocker(None)) + + image_path = folder_paths.get_annotated_filepath(image) + img = node_helpers.pillow(Image.open, image_path) + + output_images = [] + output_masks = [] + w, h = None, None + + for i in ImageSequence.Iterator(img): + i = node_helpers.pillow(ImageOps.exif_transpose, i) + + if i.mode == 'I': + i = i.point(lambda i: i * (1 / 255)) + frame = i.convert("RGB") + + if len(output_images) == 0: + w = frame.size[0] + h = frame.size[1] + + if frame.size[0] != w or frame.size[1] != h: + continue + + frame_np = np.array(frame).astype(np.float32) / 255.0 + frame_tensor = torch.from_numpy(frame_np)[None,] + if 'A' in i.getbands(): + mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + elif i.mode == 'P' and 'transparency' in i.info: + mask = np.array(i.convert('RGBA').getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + else: + mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") + output_images.append(frame_tensor) + output_masks.append(mask.unsqueeze(0)) + + if img.format == "MPO": + break + + if len(output_images) > 1: + output_image = torch.cat(output_images, dim=0) + output_mask = torch.cat(output_masks, dim=0) + else: + output_image = output_images[0] + output_mask = output_masks[0] + + return (output_image, output_mask) + + @classmethod + def IS_CHANGED(s, image, active): + if not active: + return "inactive" + image_path = folder_paths.get_annotated_filepath(image) + m = hashlib.sha256() + with open(image_path, 'rb') as f: + m.update(f.read()) + return m.digest().hex() + + @classmethod + def VALIDATE_INPUTS(s, image, active): + if not active: + return True + if not folder_paths.exists_annotated_filepath(image): + return "Invalid image file: {}".format(image) + return True + + NODE_CLASS_MAPPINGS = { "JDL_PreviewToLoad": JDL_PreviewToLoad, + "JDL_LoadImage": JDL_LoadImage, } NODE_DISPLAY_NAME_MAPPINGS = { "JDL_PreviewToLoad": "Preview to Load Image", + "JDL_LoadImage": "Load Image (Active Switch)", }