diff --git a/py/image_reel.py b/py/image_reel.py index 936b46d..c3f7742 100644 --- a/py/image_reel.py +++ b/py/image_reel.py @@ -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" -} \ No newline at end of file +} diff --git a/pyproject.toml b/pyproject.toml index 6202bf9..339d44a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"]