fix bug for ImageReelComposit node when processing Image Batch
This commit is contained in:
chflame163
+116
-64
@@ -8,6 +8,8 @@ class ImageReelPipeline:
|
||||
self.texts = {}
|
||||
self.reel_height = 0
|
||||
self.reel_border = 0
|
||||
# A Reel can contain one frame for each item in an IMAGE batch.
|
||||
self.reels = []
|
||||
|
||||
Reel = ImageReelPipeline()
|
||||
class ImageReel:
|
||||
@@ -43,35 +45,62 @@ class ImageReel:
|
||||
reel_height, border,
|
||||
image2=None, image3=None, image4=None,):
|
||||
|
||||
image_list = []
|
||||
texts = []
|
||||
for img in image1:
|
||||
i = self.resize_image_to_height(tensor2pil(img.unsqueeze(0)),reel_height)
|
||||
image_list.append(i)
|
||||
texts.append([image1_text,i.width])
|
||||
if image2 is not None:
|
||||
for img in image2:
|
||||
i = self.resize_image_to_height(tensor2pil(img.unsqueeze(0)),reel_height)
|
||||
image_list.append(i)
|
||||
texts.append([image2_text,i.width])
|
||||
if image3 is not None:
|
||||
for img in image3:
|
||||
i = self.resize_image_to_height(tensor2pil(img.unsqueeze(0)),reel_height)
|
||||
image_list.append(i)
|
||||
texts.append([image3_text,i.width])
|
||||
if image4 is not None:
|
||||
for img in image4:
|
||||
i = self.resize_image_to_height(tensor2pil(img.unsqueeze(0)),reel_height)
|
||||
image_list.append(i)
|
||||
texts.append([image4_text,i.width])
|
||||
image_batches = [self._tensor_to_pil_batch(image1, reel_height)]
|
||||
image_batches.extend([
|
||||
self._tensor_to_pil_batch(image2, reel_height),
|
||||
self._tensor_to_pil_batch(image3, reel_height),
|
||||
self._tensor_to_pil_batch(image4, reel_height),
|
||||
])
|
||||
text_labels = [image1_text, image2_text, image3_text, image4_text]
|
||||
|
||||
reel = ImageReel()
|
||||
reel.image = self.draw_reel_image(image_list, border, reel_height)
|
||||
reel.texts = texts
|
||||
batch_size = max(len(batch) for batch in image_batches)
|
||||
reel = ImageReelPipeline()
|
||||
reel.NODE_NAME = self.NODE_NAME
|
||||
for batch_index in range(batch_size):
|
||||
image_list = []
|
||||
texts = []
|
||||
for batch, text in zip(image_batches, text_labels):
|
||||
image = self._select_batch_item(batch, batch_index)
|
||||
if image is not None:
|
||||
image_list.append(image)
|
||||
texts.append([text, image.width])
|
||||
|
||||
frame = ImageReelPipeline()
|
||||
frame.image = self.draw_reel_image(image_list, border, reel_height)
|
||||
frame.texts = texts
|
||||
frame.reel_height = reel_height
|
||||
frame.reel_border = border
|
||||
reel.reels.append(frame)
|
||||
|
||||
# Keep the original single-Reel attributes for compatibility with
|
||||
# workflows or custom nodes that inspect them directly.
|
||||
if reel.reels:
|
||||
reel.image = reel.reels[0].image
|
||||
reel.texts = reel.reels[0].texts
|
||||
reel.reel_height = reel_height
|
||||
reel.reel_border = border
|
||||
return (reel,)
|
||||
|
||||
def _tensor_to_pil_batch(self, image, reel_height):
|
||||
if image is None:
|
||||
return []
|
||||
if image.dim() == 3:
|
||||
image = image.unsqueeze(0)
|
||||
if image.dim() != 4:
|
||||
raise ValueError(f"Expected an IMAGE tensor with 3 or 4 dimensions, got {image.dim()}")
|
||||
return [
|
||||
self.resize_image_to_height(tensor2pil(img.unsqueeze(0)), reel_height)
|
||||
for img in image
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _select_batch_item(batch, index):
|
||||
if not batch:
|
||||
return None
|
||||
if len(batch) == 1:
|
||||
return batch[0]
|
||||
return batch[index] if index < len(batch) else batch[-1]
|
||||
|
||||
def resize_image_to_height(self, image, target_height) -> Image:
|
||||
w = int(target_height / image.height * image.width)
|
||||
return image.resize((w, target_height), Image.LANCZOS)
|
||||
@@ -136,50 +165,73 @@ class ImageReelComposit:
|
||||
text_color = "#E5E5E5"
|
||||
|
||||
|
||||
font_space = int(font_size * 1.5)
|
||||
width = reel_1.image.width
|
||||
height = reel_1.image.height + font_space + border
|
||||
if reel_2 is not None:
|
||||
width = max(width, reel_2.image.width)
|
||||
height += reel_2.image.height + font_space + border
|
||||
if reel_3 is not None:
|
||||
width = max(width, reel_3.image.width)
|
||||
height += reel_3.image.height + font_space + border
|
||||
if reel_4 is not None:
|
||||
width = max(width, reel_4.image.width)
|
||||
height += reel_4.image.height + font_space + border
|
||||
reel_batches = [self._reel_frames(reel) for reel in (reel_1, reel_2, reel_3, reel_4)]
|
||||
batch_size = max(len(batch) for batch in reel_batches)
|
||||
for batch_index in range(batch_size):
|
||||
frames = [self._select_batch_item(batch, batch_index) for batch in reel_batches]
|
||||
frames = [frame for frame in frames if frame is not None]
|
||||
ret_images.append(pil2tensor(self._composite_frame(
|
||||
frames, font_file, font_size, border, bg_color, text_color
|
||||
)))
|
||||
|
||||
ret_image = Image.new('RGB', (width, height), color=bg_color)
|
||||
paste_y = 0
|
||||
reel1_text_image = self.draw_reel_text(reel_1, font_file, font_size, text_color)
|
||||
shadow_size = reel_1.image.height // 80
|
||||
ret_image = self.paste_drop_shadow(ret_image, reel_1.image, reel1_text_image, ((width - reel_1.image.width) // 2, paste_y),
|
||||
shadow_size, text_color)
|
||||
|
||||
paste_y += reel_1.image.height + font_space + border
|
||||
if reel_2 is not None:
|
||||
reel2_text_image = self.draw_reel_text(reel_2, font_file, font_size, text_color)
|
||||
shadow_size = reel_2.image.height // 80
|
||||
ret_image = self.paste_drop_shadow(ret_image, reel_2.image, reel2_text_image, ((width - reel_2.image.width) // 2, paste_y),
|
||||
shadow_size, text_color)
|
||||
paste_y += reel_2.image.height + font_space + border
|
||||
if reel_3 is not None:
|
||||
reel3_text_image = self.draw_reel_text(reel_3, font_file, font_size, text_color)
|
||||
shadow_size = reel_3.image.height // 80
|
||||
ret_image = self.paste_drop_shadow(ret_image, reel_3.image, reel3_text_image,((width - reel_3.image.width) // 2, paste_y),
|
||||
shadow_size, text_color)
|
||||
paste_y += reel_3.image.height + font_space + border
|
||||
if reel_4 is not None:
|
||||
reel4_text_image = self.draw_reel_text(reel_4, font_file, font_size, text_color)
|
||||
shadow_size = reel_4.image.height // 80
|
||||
ret_image = self.paste_drop_shadow(ret_image, reel_4.image, reel4_text_image,((width - reel_4.image.width) // 2, paste_y),
|
||||
shadow_size, text_color)
|
||||
|
||||
ret_images.append(pil2tensor(ret_image))
|
||||
# IMAGE batches must have a common spatial shape. Different Reel
|
||||
# inputs can produce different widths, so pad only when necessary.
|
||||
if ret_images:
|
||||
max_height = max(image.shape[1] for image in ret_images)
|
||||
max_width = max(image.shape[2] for image in ret_images)
|
||||
if any(image.shape[1] != max_height or image.shape[2] != max_width for image in ret_images):
|
||||
padded_images = []
|
||||
for image in ret_images:
|
||||
padded = torch.zeros((1, max_height, max_width, image.shape[3]), dtype=image.dtype)
|
||||
padded[:, :, :, :] = torch.tensor(tuple(int(bg_color[i:i + 2], 16) for i in (1, 3, 5)), dtype=image.dtype) / 255.0
|
||||
padded[:, :image.shape[1], :image.shape[2], :] = image
|
||||
padded_images.append(padded)
|
||||
ret_images = padded_images
|
||||
|
||||
log(f"{self.NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish')
|
||||
return (torch.cat(ret_images, dim=0),)
|
||||
|
||||
def _composite_frame(self, reels, font_file, font_size, border, bg_color, text_color):
|
||||
font_space = int(font_size * 1.5)
|
||||
width = max(reel.image.width for reel in reels)
|
||||
height = sum(reel.image.height + font_space + border for reel in reels)
|
||||
|
||||
ret_image = Image.new('RGB', (width, height), color=bg_color)
|
||||
paste_y = 0
|
||||
for reel in reels:
|
||||
reel_text_image = self.draw_reel_text(reel, font_file, font_size, text_color)
|
||||
shadow_size = reel.image.height // 80
|
||||
ret_image = self.paste_drop_shadow(
|
||||
ret_image,
|
||||
reel.image,
|
||||
reel_text_image,
|
||||
((width - reel.image.width) // 2, paste_y),
|
||||
shadow_size,
|
||||
text_color,
|
||||
)
|
||||
paste_y += reel.image.height + font_space + border
|
||||
return ret_image
|
||||
|
||||
@staticmethod
|
||||
def _reel_frames(reel):
|
||||
if reel is None:
|
||||
return []
|
||||
frames = getattr(reel, 'reels', None)
|
||||
if frames:
|
||||
return frames
|
||||
# Accept Reel objects produced by older versions of this node.
|
||||
if getattr(reel, 'image', None) is not None:
|
||||
return [reel]
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _select_batch_item(batch, index):
|
||||
if not batch:
|
||||
return None
|
||||
if len(batch) == 1:
|
||||
return batch[0]
|
||||
return batch[index] if index < len(batch) else batch[-1]
|
||||
|
||||
def paste_drop_shadow(self, background_image, image, text_image, box, shadow_size, text_color) -> Image:
|
||||
# drop shadow
|
||||
_mask = image.split()[3]
|
||||
@@ -221,4 +273,4 @@ NODE_CLASS_MAPPINGS = {
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LayerUtility: ImageReel": "LayerUtility: Image Reel",
|
||||
"LayerUtility: ImageReelComposit": "LayerUtility: Image Reel Composit"
|
||||
}
|
||||
}
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "comfyui_layerstyle"
|
||||
description = "A set of nodes for ComfyUI it generate image like Adobe Photoshop's Layer Style. the Drop Shadow is first completed node, and follow-up work is in progress."
|
||||
version = "2.0.40"
|
||||
version = "2.0.41"
|
||||
license = {text = "MIT License"}
|
||||
dependencies = ["numpy", "pillow", "torch", "matplotlib", "Scipy", "scikit_image", "scikit_learn", "opencv-contrib-python", "pymatting", "timm", "colour-science", "transformers", "blend_modes", "huggingface_hub", "loguru"]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user