image_combiner.py

"""
Image Combiner - 图片拼接节点
功能: 将两张图片水平或垂直拼接在一起，方便对比查看
"""

import torch
import numpy as np
from PIL import Image


class ImageCombiner:
    """
    图片拼接器
    将两张图片按指定方向（水平/垂直）拼接成一张图片
    """

    @classmethod
    def INPUT_TYPES(cls):
        """
        定义节点的输入参数类型
        """
        return {
            "required": {
                "image_1": ("IMAGE",),
                "image_2": ("IMAGE",),
                "direction": (["horizontal", "vertical"], {
                    "default": "horizontal"
                }),
                "gap": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 100,
                    "step": 1
                }),
                "gap_color": (["black", "white", "gray"], {
                    "default": "black"
                }),
            }
        }

    # 返回值类型定义
    RETURN_TYPES = ("IMAGE",)
    # 返回值名称定义
    RETURN_NAMES = ("combined_image",)
    # 执行函数名称
    FUNCTION = "combine_images"
    # 节点分类
    CATEGORY = "Custom/Batching"

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        """
        强制刷新机制，返回 NaN 确保每次都重新执行
        """
        return float("NaN")

    def combine_images(self, image_1, image_2, direction, gap, gap_color):
        """
        核心执行逻辑：拼接两张图片（支持批量处理）

        参数:
            image_1: 第一张图片 Tensor [B, H, W, C]
            image_2: 第二张图片 Tensor [B, H, W, C]
            direction: 拼接方向 (horizontal/vertical)
            gap: 间隙大小（像素）
            gap_color: 间隙颜色

        返回:
            (combined_image,) - 拼接后的图片批次
        """
        try:
            # 获取 batch 大小
            batch_size = image_1.shape[0]

            # 存储所有拼接后的图片
            combined_images = []

            # 确定间隙颜色
            gap_rgb = self._get_gap_color(gap_color)

            # 遍历批次中的每一对图片
            for i in range(batch_size):
                # 提取当前批次的图片
                img1 = self._tensor_to_pil(image_1[i])
                img2 = self._tensor_to_pil(image_2[i])

                # 根据方向拼接
                if direction == "horizontal":
                    combined_img = self._combine_horizontal(img1, img2, gap, gap_rgb)
                else:  # vertical
                    combined_img = self._combine_vertical(img1, img2, gap, gap_rgb)

                # 转换回 Tensor 并添加到列表
                combined_tensor = self._pil_to_tensor_single(combined_img)
                combined_images.append(combined_tensor)

            # 合并所有批次的图片
            if len(combined_images) > 0:
                result = torch.cat(combined_images, dim=0)
                return (result,)
            else:
                # 如果没有图片，返回第一张作为备用
                return (image_1,)

        except Exception as e:
            print(f"[错误] 拼接图片时发生异常: {str(e)}")
            import traceback
            traceback.print_exc()

            # 返回第一张图片作为备用
            return (image_1,)

    def _combine_horizontal(self, img1, img2, gap, gap_rgb):
        """
        水平拼接（左右排列）

        参数:
            img1: 第一张图片 (PIL.Image)
            img2: 第二张图片 (PIL.Image)
            gap: 间隙大小
            gap_rgb: 间隙颜色 RGB 元组

        返回:
            拼接后的图片 (PIL.Image)
        """
        width1, height1 = img1.size
        width2, height2 = img2.size

        # 新图片高度为两张图片中较大的高度
        new_height = max(height1, height2)
        # 新图片宽度为两张图片宽度之和加间隙
        new_width = width1 + width2 + gap

        # 创建新画布
        new_img = Image.new('RGB', (new_width, new_height), gap_rgb)

        # 粘贴第一张图片（垂直居中）
        y1 = (new_height - height1) // 2
        new_img.paste(img1, (0, y1))

        # 粘贴第二张图片（垂直居中）
        y2 = (new_height - height2) // 2
        new_img.paste(img2, (width1 + gap, y2))

        print(f"[水平拼接] {width1}x{height1} + {width2}x{height2} -> {new_width}x{new_height}")

        return new_img

    def _combine_vertical(self, img1, img2, gap, gap_rgb):
        """
        垂直拼接（上下排列）

        参数:
            img1: 第一张图片 (PIL.Image)
            img2: 第二张图片 (PIL.Image)
            gap: 间隙大小
            gap_rgb: 间隙颜色 RGB 元组

        返回:
            拼接后的图片 (PIL.Image)
        """
        width1, height1 = img1.size
        width2, height2 = img2.size

        # 新图片宽度为两张图片中较大的宽度
        new_width = max(width1, width2)
        # 新图片高度为两张图片高度之和加间隙
        new_height = height1 + height2 + gap

        # 创建新画布
        new_img = Image.new('RGB', (new_width, new_height), gap_rgb)

        # 粘贴第一张图片（水平居中）
        x1 = (new_width - width1) // 2
        new_img.paste(img1, (x1, 0))

        # 粘贴第二张图片（水平居中）
        x2 = (new_width - width2) // 2
        new_img.paste(img2, (x2, height1 + gap))

        print(f"[垂直拼接] {width1}x{height1} + {width2}x{height2} -> {new_width}x{new_height}")

        return new_img

    def _get_gap_color(self, color_name):
        """
        获取间隙颜色的 RGB 值

        参数:
            color_name: 颜色名称 (black/white/gray)

        返回:
            RGB 元组 (R, G, B)
        """
        color_map = {
            "black": (0, 0, 0),
            "white": (255, 255, 255),
            "gray": (128, 128, 128)
        }
        return color_map.get(color_name, (0, 0, 0))

    def _tensor_to_pil(self, tensor):
        """
        将 ComfyUI Tensor 转换为 PIL.Image（处理单张图片）

        参数:
            tensor: torch.Tensor [H, W, C]，值范围 0-1

        返回:
            PIL.Image 对象
        """
        # 转换为 numpy array
        img_array = tensor.cpu().numpy()

        # 反归一化到 0-255
        img_array = (img_array * 255).astype(np.uint8)

        # 转换为 PIL Image
        img = Image.fromarray(img_array, 'RGB')

        return img

    def _pil_to_tensor_single(self, img):
        """
        将 PIL.Image 转换为 ComfyUI 的 Tensor 格式（单张图片，保留 batch 维度）

        参数:
            img: PIL.Image 对象

        返回:
            torch.Tensor [1, H, W, C] 格式，值范围 0-1
        """
        # PIL.Image -> numpy array (H, W, C)
        img_array = np.array(img).astype(np.float32)

        # 归一化到 0-1
        img_array = img_array / 255.0

        # 转换为 torch.Tensor
        tensor = torch.from_numpy(img_array)

        # 增加 batch 维度: [H, W, C] -> [1, H, W, C]
        tensor = tensor.unsqueeze(0)

        return tensor