diff --git a/README.MD b/README.MD index 4a543cb..037adea 100644 --- a/README.MD +++ b/README.MD @@ -13,8 +13,8 @@ Nodes are divided into 5 groups according to their functions: LayerStyle, LayerC [中文说明点这里](./README_CN.MD) ## Update -* Comprehensive support for batch images, providing convenience for video creation. * to avoid unpredictable results, a few nodes such as CropByMask and RestoryCropBox still do not support batch images. -If it is necessary to use these nodes in the batch process, please pre convert them using the [Impact-Pack's Image Batch To Image List](https://github.com/ltdrdata/ComfyUI-Impact-Pack) node. +* All nodes have fully supported batch images, providing convenience for video creation. +(The CropByMask node only supports cuts of the same size. if a batch mask_for_crop inputted, the data from the first sheet will be used.) * Commit [RemBgUltra](#RemBgUltra) and [PixelSpread](#PixelSpread) nodes significantly improved mask quality. *RemBgUltra requires manual model download. * Commit [TextImage](#TextImage) node, it generate text images and masks. * Add new types of [blend mode](#blend) between images. now supports up to 19 blend modes. add **color_burn, color_dodge, linear_burn, linear_dodge, overlay, soft_light, hard_light, vivid_light, pin_light, linear_light** and **hard_mix**. diff --git a/README_CN.MD b/README_CN.MD index 330b8f9..4ae747d 100644 --- a/README_CN.MD +++ b/README_CN.MD @@ -10,7 +10,7 @@ * [LayerFilter](#LayerFilter)ڵṩͼЧ˾ ## ˵ -* ȫ֧ͼƬΪƵṩ㡣* ΪⲻԤĽCropByMaskRestoryCropBoxڵԲ֧ͼƬһҪʹЩڵ㣬ʹ[Impact-Pack Image Batch To Image List](https://github.com/ltdrdata/ComfyUI-Impact-Pack)ڵԤת +* нڵȫ֧ͼƬΪƵṩ㡣( CropByMask ڵ֧ͬߴг, mask_for_cropʹõһŵݡ) * [RemBgUltra](#RemBgUltra) [PixelSpread](#PixelSpread) ڵ㣬*RemBgUltraֶģ͡ * [TextImage](#TextImage) ڵ㣬ͼ֡ * ͼ֮[ģʽ](#ģʽ)ֶ֧ͣ19ֻģʽcolor_burnɫ, color_dodgeɫ, linear_burnԼ, linear_dodgeԼ, overlay, soft_light, hard_lightǿ, vivid_light, pin_light, linear_lightԹ, hard_mixʵɫϡĻģʽֻ֧ģʽĽڵ㡣 diff --git a/py/channel_shake.py b/py/channel_shake.py index 64d0da1..7babb4d 100644 --- a/py/channel_shake.py +++ b/py/channel_shake.py @@ -1,7 +1,8 @@ import math - from .imagefunc import * +NODE_NAME = 'ChannelShake' + class ChannelShake: def __init__(self): @@ -31,9 +32,9 @@ class ChannelShake: ret_images = [] - for image in image: - - _canvas = tensor2pil(image).convert('RGB') + for i in image: + i = torch.unsqueeze(i, 0) + _canvas = tensor2pil(i).convert('RGB') R, G, B = _canvas.split() x = int(math.cos(angle) * distance) y = int(math.sin(angle) * distance) @@ -53,7 +54,7 @@ class ChannelShake: ret_image = Image.merge('RGB', [R, G, B]) ret_images.append(pil2tensor(ret_image)) - log(f'ChannelShake Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/color_adapter.py b/py/color_adapter.py index 707bdc6..e3361a0 100644 --- a/py/color_adapter.py +++ b/py/color_adapter.py @@ -1,5 +1,9 @@ +import torch + from .imagefunc import * +NODE_NAME = 'ColorAdapter' + class ColorAdapter: def __init__(self): @@ -26,16 +30,26 @@ class ColorAdapter: def color_adapter(self, image, color_ref_image, opacity): ret_images = [] + # if color_ref_image.shape[0] > 0: + # color_ref_image = torch.unsqueeze(color_ref_image[0], 0) - for image in image: + l_images = [] + r_images = [] + for l in image: + l_images.append(torch.unsqueeze(l, 0)) + for r in color_ref_image: + r_images.append(torch.unsqueeze(r, 0)) + for i in range(len(l_images)): + _image = l_images[i] + _ref = r_images[i] if len(ret_images) > i else r_images[-1] - _canvas = tensor2pil(image).convert('RGB') - ret_image = color_adapter(_canvas, tensor2pil(color_ref_image).convert('RGB')) + _canvas = tensor2pil(_image).convert('RGB') + ret_image = color_adapter(_canvas, tensor2pil(_ref).convert('RGB')) ret_image = chop_image(_canvas, ret_image, blend_mode='normal', opacity=opacity) ret_images.append(pil2tensor(ret_image)) - log(f'ColorAdapter Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/color_correct_HSV.py b/py/color_correct_HSV.py index 017f2c2..8782099 100644 --- a/py/color_correct_HSV.py +++ b/py/color_correct_HSV.py @@ -1,5 +1,9 @@ +import torch + from .imagefunc import * +NODE_NAME = 'HSV' + class ColorCorrectHSV: def __init__(self): @@ -29,9 +33,9 @@ class ColorCorrectHSV: ret_images = [] - for image in image: - - _h, _s, _v = tensor2pil(image).convert('HSV').split() + for i in image: + i = torch.unsqueeze(i,0) + _h, _s, _v = tensor2pil(i).convert('HSV').split() if H != 0 : _h = image_hue_offset(_h, H) if S != 0 : @@ -42,7 +46,7 @@ class ColorCorrectHSV: ret_images.append(pil2tensor(ret_image)) - log(f'HSV Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/color_correct_LAB.py b/py/color_correct_LAB.py index f77e00a..3750915 100644 --- a/py/color_correct_LAB.py +++ b/py/color_correct_LAB.py @@ -1,5 +1,9 @@ +import torch + from .imagefunc import * +NODE_NAME = 'LAB' + class ColorCorrectLAB: def __init__(self): @@ -29,9 +33,9 @@ class ColorCorrectLAB: ret_images = [] - for image in image: - - _l, _a, _b = tensor2pil(image).convert('LAB').split() + for i in image: + i = torch.unsqueeze(i, 0) + _l, _a, _b = tensor2pil(i).convert('LAB').split() if L != 0 : _l = image_gray_offset(_l, L) if A != 0 : @@ -42,7 +46,7 @@ class ColorCorrectLAB: ret_images.append(pil2tensor(ret_image)) - log(f'LAB Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/color_correct_LUTapply.py b/py/color_correct_LUTapply.py index 708df99..862a018 100644 --- a/py/color_correct_LUTapply.py +++ b/py/color_correct_LUTapply.py @@ -1,7 +1,12 @@ import os import glob + +import torch + from .imagefunc import * +NODE_NAME = 'LUT Apply' + class ColorCorrectLUTapply: def __init__(self): @@ -27,14 +32,14 @@ class ColorCorrectLUTapply: def color_correct_LUTapply(self, image, LUT): ret_images = [] - for image in image: - - _image = tensor2pil(image) + for i in image: + i = torch.unsqueeze(i, 0) + _image = tensor2pil(i) lut_file = LUT_DICT[LUT] ret_image = lut_apply(_image, lut_file) ret_images.append(pil2tensor(ret_image)) - log(f'LUT Apply Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/color_correct_RGB.py b/py/color_correct_RGB.py index 9a67932..c25d44b 100644 --- a/py/color_correct_RGB.py +++ b/py/color_correct_RGB.py @@ -1,5 +1,9 @@ +import torch + from .imagefunc import * +NODE_NAME = 'RGB' + class ColorCorrectRGB: def __init__(self): @@ -29,9 +33,9 @@ class ColorCorrectRGB: ret_images = [] - for image in image: - - _r, _g, _b = tensor2pil(image).convert('RGB').split() + for i in image: + i = torch.unsqueeze(i,0) + _r, _g, _b = tensor2pil(i).convert('RGB').split() if R != 0 : _r = image_gray_offset(_r, R) if G != 0 : @@ -42,7 +46,7 @@ class ColorCorrectRGB: ret_images.append(pil2tensor(ret_image)) - log(f'RGB Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/color_correct_YUV.py b/py/color_correct_YUV.py index 3c0d699..44cd14b 100644 --- a/py/color_correct_YUV.py +++ b/py/color_correct_YUV.py @@ -1,5 +1,9 @@ +import torch + from .imagefunc import * +NODE_NAME = 'YUV' + class ColorCorrectYUV: def __init__(self): @@ -29,9 +33,9 @@ class ColorCorrectYUV: ret_images = [] - for image in image: - - _y, _u, _v = tensor2pil(image).convert('YCbCr').split() + for i in image: + i = torch.unsqueeze(i, 0) + _y, _u, _v = tensor2pil(i).convert('YCbCr').split() if Y != 0 : _y = image_gray_offset(_y, Y) if U != 0 : @@ -42,7 +46,7 @@ class ColorCorrectYUV: ret_images.append(pil2tensor(ret_image)) - log(f'YUV Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/color_correct_brightness&contrast.py b/py/color_correct_brightness&contrast.py index ff9b65f..6c8d35d 100644 --- a/py/color_correct_brightness&contrast.py +++ b/py/color_correct_brightness&contrast.py @@ -1,6 +1,8 @@ from PIL import ImageEnhance from .imagefunc import * +NODE_NAME = 'Brightness & Contrast' + class ColorCorrectBrightnessAndContrast: def __init__(self): @@ -30,9 +32,10 @@ class ColorCorrectBrightnessAndContrast: ret_images = [] - for image in image: + for i in image: + i = torch.unsqueeze(i,0) - _image = tensor2pil(image).convert('RGB') + _image = tensor2pil(i).convert('RGB') if brightness != 1: brightness_image = ImageEnhance.Brightness(_image) _image = brightness_image.enhance(factor=brightness) @@ -42,10 +45,9 @@ class ColorCorrectBrightnessAndContrast: if saturation != 1: color_image = ImageEnhance.Color(_image) _image = color_image.enhance(factor=saturation) - ret_images.append(pil2tensor(_image)) - log(f'Brightness & Contrast Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/color_correct_gamma.py b/py/color_correct_gamma.py index 01cecab..5075183 100644 --- a/py/color_correct_gamma.py +++ b/py/color_correct_gamma.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'Gamma' + class ColorCorrectGamma: def __init__(self): @@ -27,13 +29,13 @@ class ColorCorrectGamma: ret_images = [] - for image in image: - - ret_image = gamma_trans(tensor2pil(image), gamma) + for i in image: + i = torch.unsqueeze(i, 0) + ret_image = gamma_trans(tensor2pil(i), gamma) ret_images.append(pil2tensor(ret_image)) - log(f'Gamma Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/color_map.py b/py/color_map.py index 354ce71..686f07c 100644 --- a/py/color_map.py +++ b/py/color_map.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'ColorMap' + colormap_list = ['autumn', 'bone', 'jet', 'winter', 'rainbow', 'ocean', 'summer', 'sprint', 'cool', 'HSV', 'pink', 'hot', 'parula', 'magma', 'inferno', 'plasma', 'viridis', 'cividis', @@ -34,13 +36,15 @@ class ColorMap: ret_images = [] - for image in image: - _canvas = tensor2pil(image) + for i in image: + i = torch.unsqueeze(i, 0) + _canvas = tensor2pil(i) _image = image_to_colormap(_canvas, colormap_list.index(color_map)) ret_image = chop_image(_canvas, _image, 'normal', opacity) ret_images.append(pil2tensor(ret_image)) - log(f'ColorMap Processed {len(ret_images)} image(s).') + + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/color_overlay.py b/py/color_overlay.py index 220456d..1072df6 100644 --- a/py/color_overlay.py +++ b/py/color_overlay.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'ColorOverlay' + class ColorOverlay: def __init__(self): @@ -37,37 +39,35 @@ class ColorOverlay: l_images = [] l_masks = [] ret_images = [] - for b in background_image: - b_images.append(b) + b_images.append(torch.unsqueeze(b, 0)) for l in layer_image: - l_images.append(l) + l_images.append(torch.unsqueeze(l, 0)) m = tensor2pil(l) - if tensor2pil(l).mode == 'RGBA': - l_masks.append(m.convert('RGBA').split()[-1]) - else: - l_masks.append(Image.new('L', m.size, 'white')) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) if layer_mask is not None: l_masks = [] for m in layer_mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) - max_batch = max(len(b_images), len(l_images), len(l_masks)) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) + if len(l_masks) == 0: + log(f"Error: {NODE_NAME} skipped, because the available mask is not found.") + return (background_image,) + max_batch = max(len(b_images), len(l_images), len(l_masks)) _color = Image.new("RGB", tensor2pil(l_images[0]).size, color=color) for i in range(max_batch): background_image = b_images[i] if i < len(b_images) else b_images[-1] layer_image = l_images[i] if i < len(l_images) else l_images[-1] _mask = l_masks[i] if i < len(l_masks) else l_masks[-1] - # preprocess _canvas = tensor2pil(background_image).convert('RGB') _layer = tensor2pil(layer_image).convert('RGB') - if _mask.size != _layer.size: _mask = Image.new('L', _layer.size, 'white') - log('Warning: mask mismatch, droped!') + log(f"Warning: {NODE_NAME} mask mismatch, dropped!") # 合成layer _comp = chop_image(_layer, _color, blend_mode, opacity) @@ -75,7 +75,7 @@ class ColorOverlay: ret_images.append(pil2tensor(_canvas)) - log(f'ColorOverlay Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/crop_by_mask.py b/py/crop_by_mask.py index b38c310..cd9e8c5 100644 --- a/py/crop_by_mask.py +++ b/py/crop_by_mask.py @@ -1,5 +1,8 @@ +import torch from .imagefunc import * +NODE_NAME = 'CropByMask' + class CropByMask: def __init__(self): @@ -33,7 +36,23 @@ class CropByMask: top_reserve, bottom_reserve, left_reserve, right_reserve ): - _canvas = tensor2pil(image).convert('RGB') + ret_images = [] + ret_masks = [] + l_images = [] + l_masks = [] + for l in image: + l_images.append(torch.unsqueeze(l, 0)) + m = tensor2pil(l) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) + for m in mask_for_crop: + if invert_mask: + m = 1 - m + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) + max_batch = max(len(l_images), len(l_masks)) + # 如果有多张mask输入，使用第一张 + if mask_for_crop.shape[0] > 0: + mask_for_crop = torch.unsqueeze(mask_for_crop[0], 0) if invert_mask: mask_for_crop = 1 - mask_for_crop _mask = mask2image(mask_for_crop) @@ -46,20 +65,25 @@ class CropByMask: (x, y, width, height) = min_bounding_rect(bluredmask) if detect == "max_inscribed_rect": (x, y, width, height) = max_inscribed_rect(bluredmask) - + canvas_width, canvas_height = tensor2pil(torch.unsqueeze(image[0], 0)).convert('RGB').size x1 = x - left_reserve if x - left_reserve > 0 else 0 y1 = y - top_reserve if y - top_reserve > 0 else 0 - x2 = x + width + right_reserve if x + width + right_reserve < _canvas.width else _canvas.width - y2 = y + height + bottom_reserve if y + height + bottom_reserve < _canvas.height else _canvas.height + x2 = x + width + right_reserve if x + width + right_reserve < canvas_width else canvas_width + y2 = y + height + bottom_reserve if y + height + bottom_reserve < canvas_height else canvas_height preview_image = tensor2pil(mask_for_crop).convert('RGB') preview_image = draw_rect(preview_image, x, y, width, height, line_color="#F00000", line_width=(width+height)//100) preview_image = draw_rect(preview_image, x1, y1, x2 - x1, y2 - y1, line_color="#00F000", line_width=(width+height)//200) crop_box = (x1, y1, x2, y2) - ret_image = _canvas.crop(crop_box) - ret_mask = _mask.crop(crop_box) + for i in range(max_batch): + _canvas = tensor2pil(l_images[i]).convert('RGB') + _mask = l_masks[i] if len(l_masks) > i else l_masks[-1] + ret_images.append(pil2tensor(_canvas.crop(crop_box))) + ret_masks.append(image2mask(_mask.crop(crop_box))) + + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") + return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0), list(crop_box), pil2tensor(preview_image),) - return (pil2tensor(ret_image), image2mask(ret_mask), list(crop_box), pil2tensor(preview_image),) NODE_CLASS_MAPPINGS = { "LayerUtility: CropByMask": CropByMask diff --git a/py/drop_shadow.py b/py/drop_shadow.py index 5fdc977..5aa292c 100644 --- a/py/drop_shadow.py +++ b/py/drop_shadow.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'DropShadow' + class DropShadow: def __init__(self): @@ -42,25 +44,24 @@ class DropShadow: l_images = [] l_masks = [] ret_images = [] - for b in background_image: - b_images.append(b) + b_images.append(torch.unsqueeze(b, 0)) for l in layer_image: - l_images.append(l) + l_images.append(torch.unsqueeze(l, 0)) m = tensor2pil(l) - if tensor2pil(l).mode == 'RGBA': - l_masks.append(m.convert('RGBA').split()[-1]) - else: - l_masks.append(Image.new('L', m.size, 'white')) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) if layer_mask is not None: l_masks = [] for m in layer_mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) + if len(l_masks) == 0: + log(f"Error: {NODE_NAME} skipped, because the available mask is not found.") + return (background_image,) + max_batch = max(len(b_images), len(l_images), len(l_masks)) - - distance_x = -distance_x distance_y = -distance_y shadow_color = Image.new("RGB", tensor2pil(l_images[0]).size, color=shadow_color) @@ -70,14 +71,13 @@ class DropShadow: layer_image = l_images[i] if i < len(l_images) else l_images[-1] _mask = l_masks[i] if i < len(l_masks) else l_masks[-1] - # preprocess _canvas = tensor2pil(background_image).convert('RGB') _layer = tensor2pil(layer_image) if _mask.size != _layer.size: _mask = Image.new('L', _layer.size, 'white') - log('Warning: mask mismatch, droped!') + log(f"Warning: {NODE_NAME} mask mismatch, dropped!") if distance_x != 0 or distance_y != 0: __mask = shift_image(_mask, distance_x, distance_y) # 位移 @@ -91,7 +91,7 @@ class DropShadow: ret_images.append(pil2tensor(_canvas)) - log(f'DropShadow Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) diff --git a/py/extand_canvas.py b/py/extand_canvas.py index 73521a7..569a637 100644 --- a/py/extand_canvas.py +++ b/py/extand_canvas.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'ExtendCanvas' + class ExtendCanvas: def __init__(self): @@ -34,24 +36,47 @@ class ExtendCanvas: mask=None, ): - _image = tensor2pil(image).convert('RGB') - _mask = tensor2pil(image).convert('RGBA').split()[-1] + l_images = [] + l_masks = [] + ret_images = [] + ret_masks = [] + + for l in image: + l_images.append(torch.unsqueeze(l, 0)) + m = tensor2pil(l) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) + if mask is not None: - if invert_mask: - mask = 1 - mask - _mask = mask2image(mask).convert('L') + l_masks = [] + for m in mask: + if invert_mask: + m = 1 - m + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) + else: + l_masks.append(Image.new('L', size=tensor2pil(l_images[0]).size, color='white')) - width = _image.width + left + right - height = _image.height + top + bottom - _canvas = Image.new('RGB', (width, height), color) - _mask_canvas = Image.new('L', (width, height), "black") + max_batch = max(len(l_images), len(l_masks)) + for i in range(max_batch): + + _image = l_images[i] if i < len(l_images) else l_images[-1] + _image = tensor2pil(_image).convert('RGB') + _mask = l_masks[i] if i < len(l_masks) else l_masks[-1] + + width = _image.width + left + right + height = _image.height + top + bottom + _canvas = Image.new('RGB', (width, height), color) + _mask_canvas = Image.new('L', (width, height), "black") + + _canvas.paste(_image, box=(left,top)) + _mask_canvas.paste(_mask.convert('L'), box=(left, top)) + + ret_images.append(pil2tensor(_canvas)) + ret_masks.append(image2mask(_mask_canvas)) + + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") + return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),) - _canvas.paste(_image, box=(left,top)) - _mask_canvas.paste(_mask.convert('L'), box=(left, top)) - ret_image = _canvas - ret_mask = image2mask(_mask_canvas) - log('ExtendCanvas Processed.') - return (pil2tensor(ret_image), ret_mask,) NODE_CLASS_MAPPINGS = { "LayerUtility: ExtendCanvas": ExtendCanvas diff --git a/py/gaussian_blur.py b/py/gaussian_blur.py index 513da44..fcf77f1 100644 --- a/py/gaussian_blur.py +++ b/py/gaussian_blur.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'GaussianBlur' + class GaussianBlur: def __init__(self): @@ -27,12 +29,10 @@ class GaussianBlur: ret_images = [] - for image in image: + for i in image: + _canvas = tensor2pil(torch.unsqueeze(i, 0)).convert('RGB') - _canvas = tensor2pil(image).convert('RGB') - ret_image = gaussian_blur(_canvas, blur) - - ret_images.append(pil2tensor(ret_image)) + ret_images.append(pil2tensor(gaussian_blur(_canvas, blur))) log(f'GaussianBlur Processed {len(ret_images)} image(s).') return (torch.cat(ret_images, dim=0),) diff --git a/py/get_color_tone.py b/py/get_color_tone.py index 04244fa..894ca0f 100644 --- a/py/get_color_tone.py +++ b/py/get_color_tone.py @@ -1,3 +1,5 @@ +import torch + from .imagefunc import * class GetColorTone: @@ -24,7 +26,8 @@ class GetColorTone: OUTPUT_NODE = True def get_color_tone(self, image, mode,): - + if image.shape[0] > 0: + image = torch.unsqueeze(image[0], 0) _canvas = tensor2pil(image).convert('RGB') _canvas = gaussian_blur(_canvas, int((_canvas.width + _canvas.height) / 200)) if mode == 'main_color': diff --git a/py/get_image_size.py b/py/get_image_size.py index fb8547e..0cd0e52 100644 --- a/py/get_image_size.py +++ b/py/get_image_size.py @@ -1,3 +1,5 @@ +import torch + from .imagefunc import * class GetImageSize: @@ -24,6 +26,8 @@ class GetImageSize: def get_image_size(self, image,): + if image.shape[0] > 0: + image = torch.unsqueeze(image[0], 0) _image = tensor2pil(image) return (_image.width, _image.height, [_image.width, _image.height],) diff --git a/py/gradient_overlay.py b/py/gradient_overlay.py index 7db157b..ba91947 100644 --- a/py/gradient_overlay.py +++ b/py/gradient_overlay.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'GradientOverlay' + class GradientOverlay: def __init__(self): @@ -42,43 +44,40 @@ class GradientOverlay: l_images = [] l_masks = [] ret_images = [] - for b in background_image: - b_images.append(b) + b_images.append(torch.unsqueeze(b, 0)) for l in layer_image: - l_images.append(l) + l_images.append(torch.unsqueeze(l, 0)) m = tensor2pil(l) - if tensor2pil(l).mode == 'RGBA': - l_masks.append(m.convert('RGBA').split()[-1]) - else: - l_masks.append(Image.new('L', m.size, 'white')) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) if layer_mask is not None: l_masks = [] for m in layer_mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) + if len(l_masks) == 0: + log(f"Error: {NODE_NAME} skipped, because the available mask is not found.") + return (background_image,) max_batch = max(len(b_images), len(l_images), len(l_masks)) width, height = tensor2pil(l_images[0]).size - _gradient = gradient(start_color, end_color, width, height, float(angle)) - # start_color = RGB_to_Hex((start_alpha, start_alpha, start_alpha)) end_color = RGB_to_Hex((end_alpha, end_alpha, end_alpha)) comp_alpha = gradient(start_color, end_color, width, height, float(angle)) comp_alpha = ImageChops.invert(comp_alpha).convert('L') + for i in range(max_batch): background_image = b_images[i] if i < len(b_images) else b_images[-1] layer_image = l_images[i] if i < len(l_images) else l_images[-1] _mask = l_masks[i] if i < len(l_masks) else l_masks[-1] - # preprocess _canvas = tensor2pil(background_image).convert('RGB') _layer = tensor2pil(layer_image).convert('RGB') - if _mask.size != _layer.size: _mask = Image.new('L', _layer.size, 'white') - log('Warning: mask mismatch, droped!') + log(f"Warning: {NODE_NAME} mask mismatch, dropped!") # 合成layer _comp = chop_image(_layer, _gradient, blend_mode, opacity) @@ -88,7 +87,7 @@ class GradientOverlay: ret_images.append(pil2tensor(_canvas)) - log(f'GradientOverlay Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/image_blend.py b/py/image_blend.py index 6aeccc7..2b5fb90 100644 --- a/py/image_blend.py +++ b/py/image_blend.py @@ -1,5 +1,8 @@ +import torch from .imagefunc import * +NODE_NAME = 'ImageBlend' + class ImageBlend: def __init__(self): @@ -37,12 +40,12 @@ class ImageBlend: l_masks = [] ret_images = [] for b in background_image: - b_images.append(b) + b_images.append(torch.unsqueeze(b, 0)) for l in layer_image: - l_images.append(l) + l_images.append(torch.unsqueeze(l, 0)) m = tensor2pil(l) - if tensor2pil(l).mode == 'RGBA': - l_masks.append(m.convert('RGBA').split()[-1]) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) else: l_masks.append(Image.new('L', m.size, 'white')) if layer_mask is not None: @@ -50,7 +53,7 @@ class ImageBlend: for m in layer_mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) max_batch = max(len(b_images), len(l_images), len(l_masks)) for i in range(max_batch): background_image = b_images[i] if i < len(b_images) else b_images[-1] @@ -59,14 +62,17 @@ class ImageBlend: _canvas = tensor2pil(background_image).convert('RGB') _layer = tensor2pil(layer_image).convert('RGB') - _mask = tensor2pil(layer_image).convert('RGBA').split()[-1] + + if _mask.size != _layer.size: + _mask = Image.new('L', _layer.size, 'white') + log(f"Warning: {NODE_NAME} mask mismatch, dropped!") # 合成layer _comp = chop_image(_canvas, _layer, blend_mode, opacity) _canvas.paste(_comp, mask=_mask) ret_images.append(pil2tensor(_canvas)) - log(f'ImageBlend Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/image_blend_advance.py b/py/image_blend_advance.py index 1487903..0c6ac47 100644 --- a/py/image_blend_advance.py +++ b/py/image_blend_advance.py @@ -1,6 +1,11 @@ import copy + +import torch + from .imagefunc import * +NODE_NAME = 'ImageBlendAdvance' + class ImageBlendAdvance: def __init__(self): @@ -51,12 +56,12 @@ class ImageBlendAdvance: ret_images = [] ret_masks = [] for b in background_image: - b_images.append(b) + b_images.append(torch.unsqueeze(b, 0)) for l in layer_image: - l_images.append(l) + l_images.append(torch.unsqueeze(l, 0)) m = tensor2pil(l) - if tensor2pil(l).mode == 'RGBA': - l_masks.append(m.convert('RGBA').split()[-1]) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) else: l_masks.append(Image.new('L', m.size, 'white')) if layer_mask is not None: @@ -64,7 +69,8 @@ class ImageBlendAdvance: for m in layer_mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) + max_batch = max(len(b_images), len(l_images), len(l_masks)) for i in range(max_batch): background_image = b_images[i] if i < len(b_images) else b_images[-1] @@ -76,7 +82,7 @@ class ImageBlendAdvance: if _mask.size != _layer.size: _mask = Image.new('L', _layer.size, 'white') - log('Warning: mask mismatch, droped!') + log(f"Warning: {NODE_NAME} mask mismatch, dropped!") orig_layer_width = _layer.width orig_layer_height = _layer.height @@ -117,7 +123,7 @@ class ImageBlendAdvance: ret_images.append(pil2tensor(_canvas)) ret_masks.append(image2mask(_compmask)) - log(f'ImageBlendAdvance Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/image_channel_merge.py b/py/image_channel_merge.py index 749744a..bef5f74 100644 --- a/py/image_channel_merge.py +++ b/py/image_channel_merge.py @@ -1,5 +1,9 @@ +import torch + from .imagefunc import * +NODE_NAME = 'ImageChannelMerge' + class ImageChannelMerge: def __init__(self): @@ -28,14 +32,37 @@ class ImageChannelMerge: def image_channel_merge(self, channel_1, channel_2, channel_3, mode, channel_4=None): - _channel1 = tensor2pil(channel_1) - _channel2 = tensor2pil(channel_2) - _channel3 = tensor2pil(channel_3) - _channel4 = Image.new('L', size=_channel1.size, color='white') + c1_images = [] + c2_images = [] + c3_images = [] + c4_images = [] + ret_images = [] + + width, height = tensor2pil(torch.unsqueeze(channel_1[0], 0)).size + for c in channel_1: + c1_images.append(torch.unsqueeze(c, 0)) + for c in channel_2: + c2_images.append(torch.unsqueeze(c, 0)) + for c in channel_3: + c3_images.append(torch.unsqueeze(c, 0)) if channel_4 is not None: - _channel4 = tensor2pil(channel_4) - ret_image = image_channel_merge((_channel1, _channel2, _channel3, _channel4), mode) - return (pil2tensor(ret_image),) + for c in channel_4: + c4_images.append(torch.unsqueeze(c, 0)) + else: + c4_images.append(Image.new('L', size=(width, height), color='white')) + + max_batch = max(len(c1_images), len(c2_images), len(c3_images), len(c4_images)) + for i in range(max_batch): + c_1 = c1_images[i] if i < len(c1_images) else c1_images[-1] + c_2 = c2_images[i] if i < len(c2_images) else c2_images[-1] + c_3 = c3_images[i] if i < len(c3_images) else c3_images[-1] + c_4 = c4_images[i] if i < len(c4_images) else c4_images[-1] + ret_image = image_channel_merge((tensor2pil(c_1), tensor2pil(c_2), tensor2pil(c_3), tensor2pil(c_4)), mode) + + ret_images.append(pil2tensor(ret_image)) + + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") + return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { "LayerUtility: ImageChannelMerge": ImageChannelMerge diff --git a/py/image_channel_split.py b/py/image_channel_split.py index 7b0fb54..040c8c1 100644 --- a/py/image_channel_split.py +++ b/py/image_channel_split.py @@ -1,5 +1,9 @@ +import torch + from .imagefunc import * +NODE_NAME = 'ImageChannelSplit' + class ImageChannelSplit: def __init__(self): @@ -25,10 +29,22 @@ class ImageChannelSplit: def image_channel_split(self, image, mode): - _image = tensor2pil(image).convert('RGBA') - channel1, channel2, channel3, channel4 = image_channel_split(_image, mode) + c1_images = [] + c2_images = [] + c3_images = [] + c4_images = [] - return (pil2tensor(channel1), pil2tensor(channel2), pil2tensor(channel3), pil2tensor(channel4),) + for i in image: + i = torch.unsqueeze(i, 0) + _image = tensor2pil(i).convert('RGBA') + channel1, channel2, channel3, channel4 = image_channel_split(_image, mode) + c1_images.append(pil2tensor(channel1)) + c2_images.append(pil2tensor(channel2)) + c3_images.append(pil2tensor(channel3)) + c4_images.append(pil2tensor(channel4)) + + log(f"{NODE_NAME} Processed {len(c1_images)} image(s).") + return (torch.cat(c1_images, dim=0), torch.cat(c2_images, dim=0), torch.cat(c3_images, dim=0), torch.cat(c4_images, dim=0),) NODE_CLASS_MAPPINGS = { "LayerUtility: ImageChannelSplit": ImageChannelSplit diff --git a/py/image_mask_scale_as.py b/py/image_mask_scale_as.py index 51fa6c7..3db0d99 100644 --- a/py/image_mask_scale_as.py +++ b/py/image_mask_scale_as.py @@ -1,5 +1,8 @@ +import torch from .imagefunc import * +NODE_NAME = 'ImageMaskScaleAs' + any = AnyType("*") class ImageMaskScaleAs: @@ -59,24 +62,30 @@ class ImageMaskScaleAs: ret_masks = [] if image is not None: - for image in image: - _image = tensor2pil(image).convert('RGB') + for i in image: + i = torch.unsqueeze(i, 0) + _image = tensor2pil(i).convert('RGB') orig_width, orig_height = _image.size _image = fit_resize_image(_image, target_width, target_height, fit, resize_sampler) ret_images.append(pil2tensor(_image)) if mask is not None: - for mask in mask: - _mask = tensor2pil(mask).convert('L') + for m in mask: + m = torch.unsqueeze(m, 0) + _mask = tensor2pil(m).convert('L') orig_width, orig_height = _mask.size _mask = fit_resize_image(_mask, target_width, target_height, fit, resize_sampler).convert('L') ret_masks.append(image2mask(_mask)) if len(ret_images) > 0 and len(ret_masks) >0: + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0), [orig_width, orig_height],) elif len(ret_images) > 0 and len(ret_masks) == 0: + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0), None,) elif len(ret_images) == 0 and len(ret_masks) > 0: + log(f"{NODE_NAME} Processed {len(ret_mask)} image(s).") return (None, torch.cat(ret_masks, dim=0), [orig_width, orig_height],) else: + log(f"Error: {NODE_NAME} skipped, because the available image or mask is not found.") return (None, None,) NODE_CLASS_MAPPINGS = { diff --git a/py/image_opacity.py b/py/image_opacity.py index 4b7f889..23a9bf1 100644 --- a/py/image_opacity.py +++ b/py/image_opacity.py @@ -1,5 +1,8 @@ +import torch from .imagefunc import * +NODE_NAME = 'ImageOpacity' + class ImageOpacity: def __init__(self): @@ -29,29 +32,49 @@ class ImageOpacity: mask=None, ): - _image = tensor2pil(image).convert('RGB') - _mask = tensor2pil(image).convert('RGBA').split()[-1] + ret_images = [] + ret_masks = [] + l_images = [] + l_masks = [] + for l in image: + l_images.append(torch.unsqueeze(l, 0)) + m = tensor2pil(l) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) + else: + l_masks.append(Image.new('L', size=m.size, color='white')) if mask is not None: - _mask = mask2image(mask).convert('L') - if invert_mask: - _color = Image.new("L", _image.size, color=(255)) - else: - _color = Image.new("L", _image.size, color=(0)) + l_masks = [] + for m in mask: + if invert_mask: + m = 1 - m + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) - ret_mask = _mask - if opacity == 0: - ret_mask = _color - elif opacity < 100: - alpha = 1.0 - float(opacity) / 100 + max_batch = max(len(l_images), len(l_masks)) + + for i in range(max_batch): + _image = l_images[i] if i < len(l_images) else l_images[-1] + _image = tensor2pil(_image) + _mask = l_masks[i] if i < len(l_masks) else l_masks[-1] + if invert_mask: + _color = Image.new("L", _image.size, color=('white')) + _mask = ImageChops.invert(_mask) + else: + _color = Image.new("L", _image.size, color=('black')) + + alpha = 1 - opacity / 100.0 ret_mask = Image.blend(_mask, _color, alpha) - R, G, B, = _image.split() - if invert_mask: - ret_image = Image.merge('RGBA', (R, G, B, ImageChops.invert(ret_mask))) - else: + R, G, B, = _image.convert('RGB').split() + if invert_mask: + ret_mask = ImageChops.invert(ret_mask) ret_image = Image.merge('RGBA', (R, G, B, ret_mask)) - log('ImageOpacity Processed.') - return (pil2tensor(ret_image), image2mask(ret_mask),) + + ret_images.append(pil2tensor(ret_image)) + ret_masks.append(image2mask(ret_mask)) + + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") + return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),) NODE_CLASS_MAPPINGS = { "LayerUtility: ImageOpacity": ImageOpacity diff --git a/py/image_scale_restore.py b/py/image_scale_restore.py index 39a3b59..d3115fb 100644 --- a/py/image_scale_restore.py +++ b/py/image_scale_restore.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'ImageScaleRestore' + class ImageScaleRestore: def __init__(self): @@ -33,9 +35,24 @@ class ImageScaleRestore: mask = None, original_size = None ): - _canvas = tensor2pil(image).convert('RGB') - orig_width = _canvas.width - orig_height = _canvas.height + l_images = [] + l_masks = [] + ret_images = [] + ret_masks = [] + for l in image: + l_images.append(torch.unsqueeze(l, 0)) + m = tensor2pil(l) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) + + if mask is not None: + l_masks = [] + for m in mask: + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) + + max_batch = max(len(l_images), len(l_masks)) + + orig_width, orig_height = tensor2pil(l_images[0]).size if original_size is not None: target_width = original_size[0] target_height = original_size[1] @@ -64,13 +81,24 @@ class ImageScaleRestore: resize_sampler = Image.BOX elif method == "nearest": resize_sampler = Image.NEAREST - ret_image = _canvas.resize((target_width, target_height), resize_sampler) - ret_mask = Image.new('L', size=ret_image.size, color='white') - if mask is not None: - _mask = mask2image(mask).convert('L') - ret_mask = _mask.resize((target_width, target_height), resize_sampler) - return (pil2tensor(ret_image), image2mask(ret_mask), [orig_width, orig_height],) + for i in range(max_batch): + + _image = l_images[i] if i < len(l_images) else l_images[-1] + + _canvas = tensor2pil(_image).convert('RGB') + ret_image = _canvas.resize((target_width, target_height), resize_sampler) + ret_mask = Image.new('L', size=ret_image.size, color='white') + if mask is not None: + _mask = l_masks[i] if i < len(l_masks) else l_masks[-1] + ret_mask = _mask.resize((target_width, target_height), resize_sampler) + + ret_images.append(pil2tensor(ret_image)) + ret_masks.append(image2mask(ret_mask)) + + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") + return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0), [orig_width, orig_height],) + NODE_CLASS_MAPPINGS = { "LayerUtility: ImageScaleRestore": ImageScaleRestore diff --git a/py/image_shift.py b/py/image_shift.py index 22e9cf9..5f9f06c 100644 --- a/py/image_shift.py +++ b/py/image_shift.py @@ -1,7 +1,8 @@ import math - from .imagefunc import * +NODE_NAME = 'ImageShift' + class ImageShift: def __init__(self): @@ -40,13 +41,30 @@ class ImageShift: ret_images = [] ret_masks = [] ret_border_masks = [] - for image in image: - shift_x, shift_y = -shift_x, -shift_y - _canvas = tensor2pil(image).convert('RGB') - _mask = tensor2pil(image).convert('RGBA').split()[-1] - if mask is not None: - _mask = mask2image(mask).convert('L') + l_images = [] + l_masks = [] + + + for l in image: + l_images.append(torch.unsqueeze(l, 0)) + m = tensor2pil(l) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) + else: + l_masks.append(Image.new('L', size=m.size, color='white')) + if mask is not None: + l_masks = [] + for m in mask: + if invert_mask: + m = 1 - m + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) + + shift_x, shift_y = -shift_x, -shift_y + for i in range(len(l_images)): + _image = l_images[i] + _canvas = tensor2pil(_image).convert('RGB') + _mask = l_masks[i] if len(l_masks) < i else l_masks[-1] _border = Image.new('L', size=_canvas.size, color='black') _border = draw_border(_border, border_width=border_mask_width, color='#FFFFFF') _border = _border.resize(_canvas.size) @@ -59,7 +77,7 @@ class ImageShift: ret_masks.append(image2mask(_mask)) ret_border_masks.append(image2mask(_border)) - log(f'ImageShift Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0), torch.cat(ret_border_masks, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/inner_glow.py b/py/inner_glow.py index 53d15d6..795e409 100644 --- a/py/inner_glow.py +++ b/py/inner_glow.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'InnerGlow' + class InnerGlow: def __init__(self): @@ -44,22 +46,22 @@ class InnerGlow: l_images = [] l_masks = [] ret_images = [] - for b in background_image: - b_images.append(b) + b_images.append(torch.unsqueeze(b, 0)) for l in layer_image: - l_images.append(l) + l_images.append(torch.unsqueeze(l, 0)) m = tensor2pil(l) - if tensor2pil(l).mode == 'RGBA': - l_masks.append(m.convert('RGBA').split()[-1]) - else: - l_masks.append(Image.new('L', m.size, 'white')) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) if layer_mask is not None: l_masks = [] for m in layer_mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) + if len(l_masks) == 0: + log(f"Error: {NODE_NAME} skipped, because the available mask is not found.") + return (background_image,) max_batch = max(len(b_images), len(l_images), len(l_masks)) for i in range(max_batch): background_image = b_images[i] if i < len(b_images) else b_images[-1] @@ -71,7 +73,7 @@ class InnerGlow: if _mask.size != _layer.size: _mask = Image.new('L', _layer.size, 'white') - log('Warning: mask mismatch, droped!') + log(f"Warning: {NODE_NAME} mask mismatch, dropped!") blur_factor = blur / 20.0 grow = glow_range @@ -90,7 +92,7 @@ class InnerGlow: _layer.paste(_canvas, mask=ImageChops.invert(_mask)) ret_images.append(pil2tensor(_layer)) - log(f'InnerGlow Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) diff --git a/py/inner_shadow.py b/py/inner_shadow.py index 1562ccc..e74cb05 100644 --- a/py/inner_shadow.py +++ b/py/inner_shadow.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'InnerShadow' + class InnerShadow: def __init__(self): @@ -43,22 +45,22 @@ class InnerShadow: l_images = [] l_masks = [] ret_images = [] - for b in background_image: - b_images.append(b) + b_images.append(torch.unsqueeze(b, 0)) for l in layer_image: - l_images.append(l) + l_images.append(torch.unsqueeze(l, 0)) m = tensor2pil(l) - if tensor2pil(l).mode == 'RGBA': - l_masks.append(m.convert('RGBA').split()[-1]) - else: - l_masks.append(Image.new('L', m.size, 'white')) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) if layer_mask is not None: l_masks = [] for m in layer_mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) + if len(l_masks) == 0: + log(f"Error: {NODE_NAME} skipped, because the available mask is not found.") + return (background_image,) max_batch = max(len(b_images), len(l_images), len(l_masks)) distance_x = -distance_x distance_y = -distance_y @@ -67,14 +69,12 @@ class InnerShadow: background_image = b_images[i] if i < len(b_images) else b_images[-1] layer_image = l_images[i] if i < len(l_images) else l_images[-1] _mask = l_masks[i] if i < len(l_masks) else l_masks[-1] - # preprocess _canvas = tensor2pil(background_image).convert('RGB') _layer = tensor2pil(layer_image).convert('RGB') - if _mask.size != _layer.size: _mask = Image.new('L', _layer.size, 'white') - log('Warning: mask mismatch, droped!') + log(f"Warning: {NODE_NAME} mask mismatch, dropped!") if distance_x != 0 or distance_y != 0: __mask = shift_image(_mask, distance_x, distance_y) # 位移 @@ -88,7 +88,7 @@ class InnerShadow: ret_images.append(pil2tensor(_canvas)) - log(f'InnerShadow Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/mask_box_detect.py b/py/mask_box_detect.py index c01a281..250f570 100644 --- a/py/mask_box_detect.py +++ b/py/mask_box_detect.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'MaskBoxDetect' + class MaskBoxDetect: def __init__(self): @@ -28,6 +30,9 @@ class MaskBoxDetect: def mask_box_detect(self,mask, detect, x_adjust, y_adjust, scale_adjust): + if mask.shape[0] > 0: + mask = torch.unsqueeze(mask[0], 0) + _mask = mask2image(mask).convert('RGB') _mask = gaussian_blur(_mask, 20).convert('L') @@ -54,6 +59,7 @@ class MaskBoxDetect: preview_image = tensor2pil(mask).convert('RGB') preview_image = draw_rect(preview_image, x - x_adjust, y - y_adjust, width, height, line_color="#F00000", line_width=int(preview_image.height / 60)) preview_image = draw_rect(preview_image, x, y, width, height, line_color="#00F000", line_width=int(preview_image.height / 40)) + log(f"{NODE_NAME} Processed.") return ( pil2tensor(preview_image), round(x_percent, 2), round(y_percent, 2), _width, _height, x, y,) NODE_CLASS_MAPPINGS = { diff --git a/py/mask_edge_shrink.py b/py/mask_edge_shrink.py index 4fa3e17..af2b27e 100644 --- a/py/mask_edge_shrink.py +++ b/py/mask_edge_shrink.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'MaskEdgeShrink' + class MaskEdgeShrink: def __init__(self): @@ -36,7 +38,7 @@ class MaskEdgeShrink: for m in mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) glow_range = shrink_level * soft blur = 12 @@ -63,7 +65,7 @@ class MaskEdgeShrink: ret_masks.append(image2mask(_layer)) - log(f'MaskEdgeShrink Processed {len(ret_masks)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_masks)} image(s).") return (torch.cat(ret_masks, dim=0),) diff --git a/py/mask_gradient.py b/py/mask_gradient.py index d49e53f..076c651 100644 --- a/py/mask_gradient.py +++ b/py/mask_gradient.py @@ -1,7 +1,8 @@ import copy - from .imagefunc import * +NODE_NAME = 'MaskGradient' + class MaskGradient: def __init__(self): @@ -34,11 +35,10 @@ class MaskGradient: l_masks = [] ret_masks = [] - for m in mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) for i in range(len(l_masks)): _mask = l_masks[i] @@ -48,9 +48,7 @@ class MaskGradient: _gradient = gradient('#000000', '#FFFFFF', _mask.width, _mask.height, 0) (box_x, box_y, box_width, box_height) = min_bounding_rect(_mask) - # preview_image = mask2image(mask).convert('RGB') - # preview_image = draw_rect(preview_image, box_x, box_y, box_width, box_height, - # line_color = "#F00000", line_width = int(preview_image.height / 60)) + if gradient_side == 'top': boxsize = (width, box_height) _gradient = _gradient.transpose(Image.FLIP_TOP_BOTTOM) @@ -131,7 +129,7 @@ class MaskGradient: _canvas = chop_image(_mask, _canvas, 'normal', opacity) ret_masks.append(image2mask(_canvas)) - log(f'MaskGradient Processed {len(ret_masks)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_masks)} image(s).") return (torch.cat(ret_masks, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/mask_grow.py b/py/mask_grow.py index 7323fa5..634209a 100644 --- a/py/mask_grow.py +++ b/py/mask_grow.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = '' + class MaskGrow: def __init__(self): @@ -33,14 +35,14 @@ class MaskGrow: for m in mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) for i in range(len(l_masks)): _mask = l_masks[i] ret_masks.append(expand_mask(image2mask(_mask), grow, blur) ) - log(f'MaskGrow Processed {len(ret_masks)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_masks)} image(s).") return (torch.cat(ret_masks, dim=0),) diff --git a/py/mask_invert.py b/py/mask_invert.py index a13638c..cbe8e58 100644 --- a/py/mask_invert.py +++ b/py/mask_invert.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'MaskInvert' + class MaskInvert: def __init__(self): @@ -26,13 +28,13 @@ class MaskInvert: ret_masks = [] for m in mask: - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) for i in range(len(l_masks)): _mask = l_masks[i] ret_masks.append(mask_invert(image2mask(_mask))) - log(f'MaskInvert Processed {len(ret_masks)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_masks)} image(s).") return (torch.cat(ret_masks, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/mask_motion_blur.py b/py/mask_motion_blur.py index d30ca08..5561c2d 100644 --- a/py/mask_motion_blur.py +++ b/py/mask_motion_blur.py @@ -1,7 +1,8 @@ import copy - from .imagefunc import * +NODE_NAME = 'MaskMotionBlur' + class MaskMotionBlur: def __init__(self): @@ -35,14 +36,14 @@ class MaskMotionBlur: for m in mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) for i in range(len(l_masks)): _mask = l_masks[i] _blurimage = motion_blur(_mask, angle, blur) ret_masks.append(image2mask(_blurimage)) - log(f'MaskMotionBlur Processed {len(ret_masks)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_masks)} image(s).") return (torch.cat(ret_masks, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/mask_stroke.py b/py/mask_stroke.py index 9fe1800..c113bb0 100644 --- a/py/mask_stroke.py +++ b/py/mask_stroke.py @@ -1,6 +1,8 @@ from .imagefunc import * -class MaskStrkoe: +NODE_NAME = 'MaskStroke' + +class MaskStroke: def __init__(self): pass @@ -34,7 +36,7 @@ class MaskStrkoe: for m in mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) for i in range(len(l_masks)): _mask = l_masks[i] @@ -46,13 +48,13 @@ class MaskStrkoe: stroke_mask = subtract_mask(outer_mask, inner_mask) ret_masks.append(stroke_mask) - log(f'MaskStrkoe Processed {len(ret_masks)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_masks)} image(s).") return (torch.cat(ret_masks, dim=0),) NODE_CLASS_MAPPINGS = { - "LayerMask: MaskStrkoe": MaskStrkoe + "LayerMask: MaskStroke": MaskStroke } NODE_DISPLAY_NAME_MAPPINGS = { - "LayerMask: MaskStrkoe": "LayerMask: MaskStrkoe" + "LayerMask: MaskStroke": "LayerMask: MaskStroke" } \ No newline at end of file diff --git a/py/motion_blur.py b/py/motion_blur.py index 4664c24..b27caeb 100644 --- a/py/motion_blur.py +++ b/py/motion_blur.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'MotionBlur' + class MotionBlur: def __init__(self): @@ -28,14 +30,13 @@ class MotionBlur: ret_images = [] - for image in image: + for i in image: - _canvas = tensor2pil(image).convert('RGB') - ret_image = motion_blur(_canvas, angle, blur) + _canvas = tensor2pil(torch.unsqueeze(i, 0)).convert('RGB') - ret_images.append(pil2tensor(ret_image)) + ret_images.append(pil2tensor(motion_blur(_canvas, angle, blur))) - log(f'MotionBlur Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/outer_glow.py b/py/outer_glow.py index b61a545..ec49077 100644 --- a/py/outer_glow.py +++ b/py/outer_glow.py @@ -1,5 +1,6 @@ from .imagefunc import * +NODE_NAME = 'OuterGlow' class OuterGlow: def __init__(self): @@ -46,35 +47,33 @@ class OuterGlow: ret_images = [] for b in background_image: - b_images.append(b) + b_images.append(torch.unsqueeze(b, 0)) for l in layer_image: - l_images.append(l) + l_images.append(torch.unsqueeze(l, 0)) m = tensor2pil(l) - if tensor2pil(l).mode == 'RGBA': - l_masks.append(m.convert('RGBA').split()[-1]) - else: - l_masks.append(Image.new('L', m.size, 'white')) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) if layer_mask is not None: l_masks = [] for m in layer_mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) + if len(l_masks) == 0: + log(f"Error: {NODE_NAME} skipped, because the available mask is not found.") + return (background_image,) max_batch = max(len(b_images), len(l_images), len(l_masks)) blur_factor = blur / 20.0 for i in range(max_batch): background_image = b_images[i] if i < len(b_images) else b_images[-1] layer_image = l_images[i] if i < len(l_images) else l_images[-1] _mask = l_masks[i] if i < len(l_masks) else l_masks[-1] - # preprocess _canvas = tensor2pil(background_image).convert('RGB') _layer = tensor2pil(layer_image).convert('RGB') - if _mask.size != _layer.size: _mask = Image.new('L', _layer.size, 'white') - log('Warning: mask mismatch, droped!') - + log(f"Warning: {NODE_NAME} mask mismatch, dropped!") grow = glow_range for x in range(brightness): blur = int(grow * blur_factor) @@ -91,7 +90,7 @@ class OuterGlow: ret_images.append(pil2tensor(_canvas)) - log(f'OuterGlow Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) diff --git a/py/pixel_spread.py b/py/pixel_spread.py index 4140517..08a34b8 100644 --- a/py/pixel_spread.py +++ b/py/pixel_spread.py @@ -2,6 +2,8 @@ import copy from pymatting import * from .imagefunc import * +NODE_NAME = 'PixelSpread' + class PixelSpread: def __init__(self): @@ -34,7 +36,7 @@ class PixelSpread: ret_images = [] for l in image: - i = tensor2pil(l) + i = tensor2pil(torch.unsqueeze(l, 0)) l_images.append(i) if i.mode == 'RGBA': l_masks.append(i.split()[-1]) @@ -45,11 +47,10 @@ class PixelSpread: for m in mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) max_batch = max(len(l_images), len(l_masks)) for i in range(max_batch): - _image = l_images[i] if i < len(l_images) else l_images[-1] _mask = l_masks[i] if i < len(l_masks) else l_masks[-1] if mask_grow != 0: @@ -57,20 +58,20 @@ class PixelSpread: _mask = mask2image(_mask) i1 = pil2tensor(_image.convert('RGB')) _mask = _mask.convert('RGB') + if _image.size != _mask.size: + log(f"Error: {NODE_NAME} skipped, because the mask is not match image.") + return (image,) i_dup = copy.deepcopy(i1.cpu().numpy().astype(np.float64)) a_dup = copy.deepcopy(pil2tensor(_mask).cpu().numpy().astype(np.float64)) fg = copy.deepcopy(i1.cpu().numpy().astype(np.float64)) for index, img in enumerate(i_dup): - alpha = a_dup[index][:, :, 0] # convert to single channel - # trimap = fix_trimap(trimap, 0.01, 0.99) - # alpha = estimate_alpha_cf(image, trimap, laplacian_kwargs={"epsilon": 1e-6}, - # cg_kwargs={"maxiter": 100}) + alpha = a_dup[index][:, :, 0] fg[index], _ = estimate_foreground_ml(img, np.array(alpha), return_background=True) ret_images.append(torch.from_numpy(fg.astype(np.float32))) - log(f'PixelSpread Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/rembg_ultra.py b/py/rembg_ultra.py index 6e04732..8c7373b 100644 --- a/py/rembg_ultra.py +++ b/py/rembg_ultra.py @@ -50,6 +50,7 @@ class RemBgUltra: ret_masks = [] for i in image: + i = torch.unsqueeze(i, 0) rmbgmodel = load_model() orig_image = tensor2pil(i).convert('RGB') w,h = orig_image.size @@ -57,7 +58,6 @@ class RemBgUltra: im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1) im_tensor = torch.unsqueeze(im_tensor,0) im_tensor = torch.divide(im_tensor,255.0) - im_tensor = TF.normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0]) if torch.cuda.is_available(): im_tensor=im_tensor.cuda() @@ -70,9 +70,8 @@ class RemBgUltra: _mask = Image.fromarray(np.squeeze(im_array)).convert('L') if process_detail: # ultra edge process - i1 = torch.unsqueeze(i, dim=0) d = detail_range * 2 + 1 - i_dup = copy.deepcopy(i1.cpu().numpy().astype(np.float64)) + i_dup = copy.deepcopy(i.cpu().numpy().astype(np.float64)) a_dup = copy.deepcopy(pil2tensor(_mask.convert('RGB')).cpu().numpy().astype(np.float64)) for index, img in enumerate(i_dup): trimap = a_dup[index][:,:,0] # convert to single channel diff --git a/py/restore_crop_box.py b/py/restore_crop_box.py index bcb67a7..7dd1599 100644 --- a/py/restore_crop_box.py +++ b/py/restore_crop_box.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'RestoreCropBox' + class RestoreCropBox: def __init__(self): @@ -30,18 +32,45 @@ class RestoreCropBox: croped_mask=None ): - _canvas = tensor2pil(background_image).convert('RGB') - _layer = tensor2pil(croped_image).convert('RGB') - _mask = Image.new('L', size=_layer.size, color='white') + b_images = [] + l_images = [] + l_masks = [] + ret_images = [] + ret_masks = [] + for b in background_image: + b_images.append(torch.unsqueeze(b, 0)) + for l in croped_image: + l_images.append(torch.unsqueeze(l, 0)) + m = tensor2pil(l) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) + else: + l_masks.append(Image.new('L', size=m.size, color='white')) if croped_mask is not None: - if invert_mask: - croped_mask = 1 - croped_mask - _mask = mask2image(croped_mask).convert('L') - ret_mask = Image.new('L', size=_canvas.size, color='black') - _canvas.paste(_layer, box=tuple(crop_box), mask=_mask) - ret_mask.paste(_mask, box=tuple(crop_box)) + l_masks = [] + for m in croped_mask: + if invert_mask: + m = 1 - m + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) + + max_batch = max(len(b_images), len(l_images), len(l_masks)) + for i in range(max_batch): + background_image = b_images[i] if i < len(b_images) else b_images[-1] + croped_image = l_images[i] if i < len(l_images) else l_images[-1] + _mask = l_masks[i] if i < len(l_masks) else l_masks[-1] + + _canvas = tensor2pil(background_image).convert('RGB') + _layer = tensor2pil(croped_image).convert('RGB') + + ret_mask = Image.new('L', size=_canvas.size, color='black') + _canvas.paste(_layer, box=tuple(crop_box), mask=_mask) + ret_mask.paste(_mask, box=tuple(crop_box)) + ret_images.append(pil2tensor(_canvas)) + ret_masks.append(image2mask(ret_mask)) + + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") + return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),) - return (pil2tensor(_canvas), image2mask(ret_mask),) NODE_CLASS_MAPPINGS = { "LayerUtility: RestoreCropBox": RestoreCropBox diff --git a/py/skin_beauty.py b/py/skin_beauty.py index 2d3c3e4..8d777ac 100644 --- a/py/skin_beauty.py +++ b/py/skin_beauty.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'SkinBeauty' + class SkinBeauty: def __init__(self): @@ -30,9 +32,9 @@ class SkinBeauty: ret_images = [] ret_masks = [] - for image in image: - - _canvas = tensor2pil(image).convert('RGB') + for i in image: + i = torch.unsqueeze(i, 0) + _canvas = tensor2pil(i).convert('RGB') _R, _, _, _ = image_channel_split(_canvas, mode='RGB') _otsumask = gray_threshold(_R, otsu=True) _removebkgd = remove_background(_R, _otsumask, '#000000') @@ -48,7 +50,7 @@ class SkinBeauty: ret_images.append(pil2tensor(_canvas)) ret_masks.append(image2mask(light_mask)) - log(f'SkinBeauty Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),) diff --git a/py/soft_light.py b/py/soft_light.py index fd37254..a6e1c18 100644 --- a/py/soft_light.py +++ b/py/soft_light.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'SoftLight' + class SoftLight: def __init__(self): @@ -29,10 +31,10 @@ class SoftLight: ret_images = [] - for image in image: - + for i in image: + i = torch.unsqueeze(i, 0) blend_mode = 'screen' - _canvas = tensor2pil(image).convert('RGB') + _canvas = tensor2pil(i).convert('RGB') blur = int((_canvas.width + _canvas.height) / 200 * soft) _otsumask = gray_threshold(_canvas, otsu=True) _removebkgd = remove_background(_canvas, _otsumask, '#000000').convert('L') @@ -47,7 +49,7 @@ class SoftLight: ret_images.append(pil2tensor(_canvas)) - log(f'SoftLight Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/stroke.py b/py/stroke.py index a689a77..3d02301 100644 --- a/py/stroke.py +++ b/py/stroke.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'Storke' + class Stroke: def __init__(self): @@ -42,20 +44,22 @@ class Stroke: l_masks = [] ret_images = [] for b in background_image: - b_images.append(b) + b_images.append(torch.unsqueeze(b, 0)) for l in layer_image: - l_images.append(l) + l_images.append(torch.unsqueeze(l, 0)) m = tensor2pil(l) - if tensor2pil(l).mode == 'RGBA': - l_masks.append(m.convert('RGBA').split()[-1]) - else: - l_masks.append(Image.new('L', m.size, 'white')) + if m.mode == 'RGBA': + l_masks.append(m.split()[-1]) if layer_mask is not None: l_masks = [] for m in layer_mask: if invert_mask: m = 1 - m - l_masks.append(tensor2pil(m).convert('L')) + l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) + if len(l_masks) == 0: + log(f"Error: {NODE_NAME} skipped, because the available mask is not found.") + return (background_image,) + max_batch = max(len(b_images), len(l_images), len(l_masks)) grow_offset = int(stroke_width / 2) @@ -72,7 +76,7 @@ class Stroke: if _mask.size != _layer.size: _mask = Image.new('L', _layer.size, 'white') - log('Warning: mask mismatch, droped!') + log(f"Warning: {NODE_NAME} mask mismatch, dropped!") inner_mask = expand_mask(image2mask(_mask), inner_stroke, blur) outer_mask = expand_mask(image2mask(_mask), outer_stroke, blur) @@ -84,7 +88,7 @@ class Stroke: ret_images.append(pil2tensor(_canvas)) - log(f'Stroke Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = { diff --git a/py/text_image.py b/py/text_image.py index 25260d6..da6c247 100644 --- a/py/text_image.py +++ b/py/text_image.py @@ -3,6 +3,7 @@ import random from PIL import ImageFont from .imagefunc import * +NODE_NAME = 'TextImage' any = AnyType("*") class TextImage: @@ -127,6 +128,7 @@ class TextImage: _color = Image.new('RGB', size=(width, height), color=text_color) _canvas.paste(_color, mask=_mask.convert('L')) _canvas = RGB2RGBA(_canvas, _mask) + log(f"{NODE_NAME} Processed.") return (pil2tensor(_canvas), image2mask(_mask),) NODE_CLASS_MAPPINGS = { diff --git a/py/water_color.py b/py/water_color.py index 0d92569..1cfe058 100644 --- a/py/water_color.py +++ b/py/water_color.py @@ -1,5 +1,7 @@ from .imagefunc import * +NODE_NAME = 'WaterColor' + class WaterColor: def __init__(self): @@ -29,15 +31,15 @@ class WaterColor: ret_images = [] - for image in image: - - _canvas = tensor2pil(image).convert('RGB') + for i in image: + i = torch.unsqueeze(i, 0) + _canvas = tensor2pil(i).convert('RGB') _image = image_watercolor(_canvas, level=101-line_density) ret_image = chop_image(_canvas, _image, 'normal', opacity) ret_images.append(pil2tensor(ret_image)) - log(f'WaterColor Processed {len(ret_images)} image(s).') + log(f"{NODE_NAME} Processed {len(ret_images)} image(s).") return (torch.cat(ret_images, dim=0),) NODE_CLASS_MAPPINGS = {