From 9ffd969b05ce37d6496d6dcd2d307aabf93650d5 Mon Sep 17 00:00:00 2001 From: BigWhiteFly <35857821+BigWhiteFly@users.noreply.github.com> Date: Mon, 19 May 2025 17:17:31 +0800 Subject: [PATCH] initialize project initialize project --- README.md | 2 - __init__.py | 2 + nodes.py | 202 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 204 insertions(+), 2 deletions(-) create mode 100644 __init__.py create mode 100644 nodes.py diff --git a/README.md b/README.md index cf7c75e..e69de29 100644 --- a/README.md +++ b/README.md @@ -1,2 +0,0 @@ -# ComfyUI-ImageConcat -concatenate all images in floders, concatenate caption txt files for trainning loras. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..8e2bed7 --- /dev/null +++ b/__init__.py @@ -0,0 +1,2 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..8ba24b1 --- /dev/null +++ b/nodes.py @@ -0,0 +1,202 @@ +import os +import torch +import numpy as np +from PIL import Image +import PIL.ImageOps as ImageOps +import folder_paths +import comfy.utils + +class ImageConcatenateBatchWithTxt: + def __init__(self): + pass + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image_folder1": ("STRING", {"default": ""}), + "image_folder2": ("STRING", {"default": ""}), + "output_folder": ("STRING", {"default": ""}), + "prompt_prefix": ("STRING", {"default": ""}), + "prompt_subfix": ("STRING", {"default": ""}), + } + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "process_images" + CATEGORY = "image" + + def concanate(self, image1, image2, direction, match_image_size, first_image_shape=None): + # Check if the batch sizes are different + batch_size1 = image1.shape[0] + batch_size2 = image2.shape[0] + + if batch_size1 != batch_size2: + # Calculate the number of repetitions needed + max_batch_size = max(batch_size1, batch_size2) + repeats1 = max_batch_size // batch_size1 + repeats2 = max_batch_size // batch_size2 + + # Repeat the images to match the largest batch size + image1 = image1.repeat(repeats1, 1, 1, 1) + image2 = image2.repeat(repeats2, 1, 1, 1) + + if match_image_size: + # Use first_image_shape if provided; otherwise, default to image1's shape + target_shape = first_image_shape if first_image_shape is not None else image1.shape + + original_height = image2.shape[1] + original_width = image2.shape[2] + original_aspect_ratio = original_width / original_height + + if direction in ['left', 'right']: + # Match the height and adjust the width to preserve aspect ratio + target_height = target_shape[1] # B, H, W, C format + target_width = int(target_height * original_aspect_ratio) + elif direction in ['up', 'down']: + # Match the width and adjust the height to preserve aspect ratio + target_width = target_shape[2] # B, H, W, C format + target_height = int(target_width / original_aspect_ratio) + + # Adjust image2 to the expected format for common_upscale + image2_for_upscale = image2.movedim(-1, 1) # Move C to the second position (B, C, H, W) + + # Resize image2 to match the target size while preserving aspect ratio + image2_resized = comfy.utils.common_upscale(image2_for_upscale, target_width, target_height, "lanczos", "disabled") + + # Adjust image2 back to the original format (B, H, W, C) after resizing + image2_resized = image2_resized.movedim(1, -1) + else: + image2_resized = image2 + + # Concatenate based on the specified direction + if direction == 'right': + concatenated_image = torch.cat((image1, image2_resized), dim=2) # Concatenate along width + elif direction == 'down': + concatenated_image = torch.cat((image1, image2_resized), dim=1) # Concatenate along height + elif direction == 'left': + concatenated_image = torch.cat((image2_resized, image1), dim=2) # Concatenate along width + elif direction == 'up': + concatenated_image = torch.cat((image2_resized, image1), dim=1) # Concatenate along height + return concatenated_image, + + def load_image(self, image_path): + # image_path = os.path.join(image_folder1,image_file_name) + + img = Image.open(image_path) + + output_images = [] + output_masks = [] + w, h = None, None + + excluded_formats = ['MPO'] + + i= img + i = 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: + return None + + 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) + 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: + output_image = output_images[0] + output_mask = output_masks[0] + return output_image + #return (output_image, output_mask) + + def process_images(self, image_folder1, image_folder2, output_folder,prompt_prefix, prompt_subfix): + + if not os.path.isabs(output_folder): + output_folder = os.path.join(folder_paths.base_path, os.path.normpath(output_folder)) + if not os.path.isabs(image_folder1): + image_folder1 = os.path.join(folder_paths.base_path, os.path.normpath(image_folder1)) + if not os.path.isabs(image_folder2): + image_folder2 = os.path.join(folder_paths.base_path, os.path.normpath(image_folder2)) + + if not os.path.exists(image_folder1) or not os.path.exists(image_folder2): + print("Error: One or both image folders do not exist.") + print(f"Image folder 1: {image_folder1}") + print(f"Image folder 2: {image_folder2}") + print(f"Output folder: {output_folder}") + return image_folder1, image_folder2, output_folder + + # 确保输出文件夹存在 + os.makedirs(output_folder, exist_ok=True) + + # 获取两个文件夹中的图片文件 + image_extensions = ('.png', '.jpg', '.jpeg') + images1 = [f for f in os.listdir(image_folder1) if f.lower().endswith(image_extensions)] + images2 = [f for f in os.listdir(image_folder2) if f.lower().endswith(image_extensions)] + + + for img1_name in images1: + img1_path = os.path.join(image_folder1, img1_name) + txt1_name = os.path.splitext(img1_name)[0] + '.txt' + txt1_path = os.path.join(image_folder1, txt1_name) + + # 加载第一张图片 + img1 = self.load_image(img1_path) + + for img2_name in images2: + img2_path = os.path.join(image_folder2, img2_name) + txt2_name = os.path.splitext(img2_name)[0] + '.txt' + txt2_path = os.path.join(image_folder2, txt2_name) + + # 加载第二张图片 + img2 = self.load_image(img2_path) + + # 调用 concanate 函数 + concatenated_image, = self.concanate(img1, img2, 'right', True) + + # 保存拼接后的图片 + output_img_name = os.path.splitext(img1_name)[0] + '_' + os.path.splitext(img2_name)[0] + '.png' + output_img_path = os.path.join(output_folder, output_img_name) + output_image = concatenated_image[0].cpu().numpy()#concatenated_image[0].movedim(0, -1).cpu().numpy() + output_image = np.clip(255. * output_image, 0, 255).astype(np.uint8) + Image.fromarray(output_image).save(output_img_path) + + # 读取并合并 txt 文件 + try: + with open(txt1_path, 'r', encoding='utf-8') as f1, open(txt2_path, 'r', encoding='utf-8') as f2: + txt1_content = f1.read() + txt2_content = f2.read() + combined_content = prompt_prefix+ ',' + txt1_content + ',' + txt2_content+ ',' + prompt_subfix + + # 保存合并后的 txt 文件 + output_txt_name = os.path.splitext(txt1_name)[0] + '_' + os.path.splitext(txt2_name)[0] + '.txt' + output_txt_path = os.path.join(output_folder, output_txt_name) + with open(output_txt_path, 'w', encoding='utf-8') as f_out: + f_out.write(combined_content) + except FileNotFoundError: + continue + + return output_folder, + + +NODE_CLASS_MAPPINGS = { + "ImageConcatenateBatchWithTxt": ImageConcatenateBatchWithTxt, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageConcatenateBatchWithTxt": "Image Concatenate Batch With Txt", +} +