dual_loader.py

"""
Dual Folder Filename Matcher - 核心逻辑实现
功能: 从两个文件夹中按文件名匹配加载图片对
"""

import os
from pathlib import Path
import torch
import numpy as np
from PIL import Image, ImageOps


class DualFolderFilenameMatcher:
    """
    双文件夹文件名匹配加载器
    从两个文件夹中根据文件名（不含扩展名）匹配图片并成对输出
    """

    # 类变量：用于自动递增模式
    _auto_counters = {}  # {folder_pair_key: current_index}

    @classmethod
    def INPUT_TYPES(cls):
        """
        定义节点的输入参数类型
        """
        return {
            "required": {
                "folder_path_A": ("STRING", {
                    "default": "",
                    "multiline": False
                }),
                "folder_path_B": ("STRING", {
                    "default": "",
                    "multiline": False
                }),
                "batch_index": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 99999,
                    "step": 1
                }),
                "auto_mode": (["manual", "sequential", "loop"], {
                    "default": "manual"
                })
            }
        }

    # 返回值类型定义
    RETURN_TYPES = ("IMAGE", "IMAGE", "STRING", "INT")
    # 返回值名称定义
    RETURN_NAMES = ("image_A", "image_B", "filename_stem", "total_matches")
    # 执行函数名称
    FUNCTION = "load_images"
    # 节点分类
    CATEGORY = "Custom/Batching"

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        """
        强制刷新机制，返回 NaN 确保每次都重新执行，不使用缓存
        这对批量处理至关重要
        """
        return float("NaN")

    def load_images(self, folder_path_A, folder_path_B, batch_index, auto_mode):
        """
        核心执行逻辑：加载匹配的图片对

        参数:
            folder_path_A: 文件夹A的路径
            folder_path_B: 文件夹B的路径
            batch_index: 批次索引（manual 模式下使用）
            auto_mode: 自动模式 (manual/sequential/loop)

        返回:
            (image_A, image_B, filename_stem, total_matches)
        """
        # 支持的图片格式（忽略大小写）
        supported_formats = {'.png', '.jpg', '.jpeg', '.webp', '.bmp', '.tiff', '.gif'}

        try:
            # 1. 检查文件夹是否存在
            path_A = Path(folder_path_A)
            path_B = Path(folder_path_B)

            if not path_A.exists() or not path_A.is_dir():
                print(f"[错误] 文件夹 A 不存在: {folder_path_A}")
                return self._return_blank_images("folder_A_not_found")

            if not path_B.exists() or not path_B.is_dir():
                print(f"[错误] 文件夹 B 不存在: {folder_path_B}")
                return self._return_blank_images("folder_B_not_found")

            # 2. 扫描文件夹并构建文件名映射字典
            dict_A = {}  # {文件名stem: 完整文件名}
            dict_B = {}

            # 扫描文件夹 A
            for file in path_A.iterdir():
                if file.is_file() and file.suffix.lower() in supported_formats:
                    stem = file.stem  # 文件名（不含扩展名）
                    dict_A[stem] = file.name

            # 扫描文件夹 B
            for file in path_B.iterdir():
                if file.is_file() and file.suffix.lower() in supported_formats:
                    stem = file.stem
                    dict_B[stem] = file.name

            # 3. 找出两个文件夹中文件名的交集
            common_stems = set(dict_A.keys()) & set(dict_B.keys())

            if not common_stems:
                print(f"[警告] 在两个文件夹中没有找到匹配的文件名")
                print(f"  文件夹 A: {len(dict_A)} 个图片文件")
                print(f"  文件夹 B: {len(dict_B)} 个图片文件")
                print(f"  交集: 0 个匹配")
                return self._return_blank_images("no_matches")

            # 4. 对交集结果进行排序（确保每次运行顺序一致）
            sorted_stems = sorted(common_stems)
            total_matches = len(sorted_stems)

            print(f"[信息] 找到 {total_matches} 对匹配的图片文件")

            # 5. 根据 auto_mode 确定实际使用的索引
            folder_pair_key = f"{folder_path_A}||{folder_path_B}"

            if auto_mode == "manual":
                # 手动模式：使用用户提供的 batch_index（取模防止越界）
                current_index = batch_index % total_matches
                print(f"[模式] 手动模式 - 使用 batch_index={batch_index}")

            elif auto_mode == "sequential":
                # 顺序模式：自动递增，不循环
                if folder_pair_key not in self._auto_counters:
                    self._auto_counters[folder_pair_key] = 0

                current_index = self._auto_counters[folder_pair_key]

                # 如果超出范围，停在最后一张
                if current_index >= total_matches:
                    current_index = total_matches - 1
                    print(f"[模式] 顺序模式 - 已达到最后一张，停在索引 {current_index}")
                else:
                    print(f"[模式] 顺序模式 - 当前索引 {current_index}，下次将自动递增")
                    # 递增计数器（下次执行时使用）
                    self._auto_counters[folder_pair_key] += 1

            else:  # loop
                # 循环模式：自动递增，循环回到开头
                if folder_pair_key not in self._auto_counters:
                    self._auto_counters[folder_pair_key] = 0

                current_index = self._auto_counters[folder_pair_key] % total_matches
                print(f"[模式] 循环模式 - 当前索引 {current_index}")

                # 递增计数器（下次执行时使用）
                self._auto_counters[folder_pair_key] += 1

            current_stem = sorted_stems[current_index]

            # 6. 获取完整文件路径
            file_A = path_A / dict_A[current_stem]
            file_B = path_B / dict_B[current_stem]

            print(f"[当前] 加载匹配对 ({current_index + 1}/{total_matches}): {current_stem}")
            print(f"  文件 A: {file_A.name}")
            print(f"  文件 B: {file_B.name}")

            # 7. 加载图片
            img_A = self._load_image(file_A)
            img_B = self._load_image(file_B)

            # 8. 统一图片尺寸（长边对齐到较小值）
            img_A, img_B = self._resize_to_match(img_A, img_B)

            # 9. 转换为 ComfyUI 的 Tensor 格式
            tensor_A = self._pil_to_tensor(img_A)
            tensor_B = self._pil_to_tensor(img_B)

            return (tensor_A, tensor_B, current_stem, total_matches)

        except Exception as e:
            print(f"[错误] 加载图片时发生异常: {str(e)}")
            import traceback
            traceback.print_exc()
            return self._return_blank_images("exception")

    def _load_image(self, file_path):
        """
        加载单张图片并进行预处理

        参数:
            file_path: 图片文件路径

        返回:
            PIL.Image 对象 (RGB格式)
        """
        img = Image.open(file_path)

        # 处理 EXIF 旋转信息
        img = ImageOps.exif_transpose(img)

        # 转换为 RGB（确保格式一致）
        if img.mode != 'RGB':
            img = img.convert('RGB')

        return img

    def _resize_to_match(self, img_A, img_B):
        """
        统一两张图片的尺寸（长边对齐到较小值）

        参数:
            img_A: 图片 A (PIL.Image)
            img_B: 图片 B (PIL.Image)

        返回:
            (调整后的 img_A, 调整后的 img_B)
        """
        # 获取两张图片的尺寸
        width_A, height_A = img_A.size
        width_B, height_B = img_B.size

        # 计算长边
        max_dim_A = max(width_A, height_A)
        max_dim_B = max(width_B, height_B)

        # 如果尺寸已经相同，直接返回
        if width_A == width_B and height_A == height_B:
            return img_A, img_B

        # 找到较小的长边作为目标尺寸
        target_max_dim = min(max_dim_A, max_dim_B)

        # 缩放图片 A（如果需要）
        if max_dim_A > target_max_dim:
            scale = target_max_dim / max_dim_A
            new_width_A = int(width_A * scale)
            new_height_A = int(height_A * scale)
            img_A = img_A.resize((new_width_A, new_height_A), Image.LANCZOS)
            print(f"  图片 A 缩放: {width_A}x{height_A} -> {new_width_A}x{new_height_A}")

        # 缩放图片 B（如果需要）
        if max_dim_B > target_max_dim:
            scale = target_max_dim / max_dim_B
            new_width_B = int(width_B * scale)
            new_height_B = int(height_B * scale)
            img_B = img_B.resize((new_width_B, new_height_B), Image.LANCZOS)
            print(f"  图片 B 缩放: {width_B}x{height_B} -> {new_width_B}x{new_height_B}")

        return img_A, img_B

    def _pil_to_tensor(self, img):
        """
        将 PIL.Image 转换为 ComfyUI 的 Tensor 格式

        参数:
            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

    def _return_blank_images(self, reason):
        """
        返回黑色空白图（异常情况下使用，避免工作流崩溃）

        参数:
            reason: 返回空白图的原因

        返回:
            (blank_image, blank_image, reason, 0)
        """
        # 创建 512x512 黑色图片
        blank_img = Image.new('RGB', (512, 512), (0, 0, 0))
        blank_tensor = self._pil_to_tensor(blank_img)

        return (blank_tensor, blank_tensor, reason, 0)