diff --git a/PainterNode/painter_node.py b/PainterNode/painter_node.py index 7e12c88..5790969 100644 --- a/PainterNode/painter_node.py +++ b/PainterNode/painter_node.py @@ -6,11 +6,13 @@ from aiohttp import web import base64 from io import BytesIO import time -from PIL import Image, ImageOps +from PIL import Image, ImageOps, ImageSequence import torch import numpy as np import glob + import folder_paths +import node_helpers # Directory node save settings CHUNK_SIZE = 1024 @@ -402,18 +404,50 @@ class PainterNode(object): # end - Piping image input image_path = folder_paths.get_annotated_filepath(image) + + img = node_helpers.pillow(Image.open, image_path) - i = Image.open(image_path) - i = ImageOps.exif_transpose(i) - image = i.convert("RGB") - image = np.array(image).astype(np.float32) / 255.0 - image = torch.from_numpy(image)[None,] - if "A" in i.getbands(): - mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0 - mask = 1.0 - torch.from_numpy(mask) + output_images = [] + output_masks = [] + w, h = None, None + + excluded_formats = ['MPO'] + + 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)) + image = i.convert("RGB") + + if len(output_images) == 0: + w = image.size[0] + h = image.size[1] + + if image.size[0] != w or image.size[1] != h: + continue + + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[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(image) + output_masks.append(mask.unsqueeze(0)) + + if len(output_images) > 1 and img.format not in excluded_formats: + output_image = torch.cat(output_images, dim=0) + output_mask = torch.cat(output_masks, dim=0) else: - mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") - return (image, mask.unsqueeze(0)) + output_image = output_images[0] + output_mask = output_masks[0] + + return (output_image, output_mask) @classmethod def IS_CHANGED(self, image, unique_id, update_node=True, images=None): diff --git a/PoseNode/pose_node.py b/PoseNode/pose_node.py index 405bfc5..23e90ba 100644 --- a/PoseNode/pose_node.py +++ b/PoseNode/pose_node.py @@ -1,10 +1,11 @@ import hashlib import os -from PIL import Image, ImageOps +from PIL import Image, ImageOps, ImageSequence import torch import numpy as np -import folder_paths +import folder_paths +import node_helpers class PoseNode(object): @classmethod @@ -32,21 +33,49 @@ class PoseNode(object): def output_pose(self, image): image_path = folder_paths.get_annotated_filepath(image) - i = Image.open(image_path) - i = ImageOps.exif_transpose(i) + img = node_helpers.pillow(Image.open, image_path) - if i.mode == 'RGBA': - mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0 - mask = 1.0 - torch.from_numpy(mask) + output_images = [] + output_masks = [] + w, h = None, None + + excluded_formats = ['MPO'] + + 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)) image = i.convert("RGB") + + if len(output_images) == 0: + w = image.size[0] + h = image.size[1] + + if image.size[0] != w or image.size[1] != h: + continue + + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[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(image) + output_masks.append(mask.unsqueeze(0)) + + if len(output_images) > 1 and img.format not in excluded_formats: + output_image = torch.cat(output_images, dim=0) + output_mask = torch.cat(output_masks, dim=0) else: - mask = torch.zeros((i.height, i.width), dtype=torch.float32, device="cpu") - image = i.convert("RGB") + output_image = output_images[0] + output_mask = output_masks[0] - image = np.array(image).astype(np.float32) / 255.0 - image = torch.from_numpy(image)[None,] - - return (image, mask.unsqueeze(0)) + return (output_image, output_mask) @classmethod def IS_CHANGED(self, image): diff --git a/pyproject.toml b/pyproject.toml index 6e47f4c..2b12a86 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui_custom_nodes_alekpet" description = "Nodes: PoseNode, PainterNode, TranslateTextNode, TranslateCLIPTextEncodeNode, DeepTranslatorTextNode, DeepTranslatorCLIPTextEncodeNode, ArgosTranslateTextNode, ArgosTranslateCLIPTextEncodeNode, ChatGLM4TranslateCLIPTextEncodeNode, ChatGLM4TranslateTextNode, ChatGLM4InstructNode, ChatGLM4InstructMediaNode, PreviewTextNode, HexToHueNode, ColorsCorrectNode, IDENode." -version = "1.0.83" +version = "1.0.84" license = { file = "LICENSE" } [project.urls]