Charlweed-ComfyUI-ZipBatchLoader/zip_batch_loader.py

"""
Core logic for the ComfyUI ZIP Batch Loader custom node.
"""

import io
import logging
import os
import zipfile

import numpy as np
import torch  # pylint: disable=import-error
from PIL import Image

import folder_paths  # pylint: disable=import-error


# noinspection PyPep8Naming
class ZipBatchLoader:
    """
    A custom ComfyUI node that loads a batch of images and masks from a ZIP archive.

    This node reads a ZIP file from the ComfyUI input directory, extracts image files,
    sorts them alphabetically, and stacks them into batched tensors. It is designed
    to handle standardized datasets (like Geomancy bundles) efficiently in-memory.

    Attributes:
        RETURN_TYPES (tuple): The types of data returned by the node (IMAGE, MASK, INT).
        RETURN_NAMES (tuple): The names of the returned outputs (image, mask, count).
        FUNCTION (str): The name of the entry point method (load_from_zip).
        CATEGORY (str): The category under which the node appears in ComfyUI.
    """

    @classmethod
    def INPUT_TYPES(cls):  # pylint: disable=invalid-name
        """
        Defines the input parameters for the node.

        Queries the ComfyUI input directory for available ZIP files to populate
        the dropdown menu.

        Returns:
            dict: A dictionary defining the required and optional inputs.
                - zip_file: A dropdown list of .zip files in the input folder.
                - heterogeneous_dimensions: A boolean toggle. If True, allows images
                  of different sizes to be skipped instead of raising an error.
        """
        files = []
        try:
            input_dir = folder_paths.get_input_directory()
            files = [
                f
                for f in os.listdir(input_dir)
                if os.path.isfile(os.path.join(input_dir, f)) and f.lower().endswith(".zip")
            ]
        except OSError as e:
            logging.error("[ZipBatchLoader] Failed to read input directory: %s", e)

        return {
            "required": {
                "zip_file": (files,),
                "heterogeneous_dimensions": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK", "INT")
    RETURN_NAMES = ("image", "mask", "count")
    FUNCTION = "load_from_zip"
    CATEGORY = "image"

    def _process_image_file(self, archive, filename, first_dims, heterogeneous_dims):
        """
        Extracts and processes a single image file from the ZIP archive.

        This helper method reads the image data into memory, validates its dimensions
        against the first image in the batch, and converts it into PyTorch tensors.

        Args:
            archive (zipfile.ZipFile): The open ZIP archive object.
            filename (str): The internal path/name of the file to process.
            first_dims (tuple, optional): The (width, height) of the first image
                processed in this batch. Used for dimension validation.
            heterogeneous_dims (bool): If True, skips images with mismatched dimensions.

        Returns:
            tuple: A tuple containing:
                - rgb_tensor (torch.Tensor or None): Normalized RGB image tensor [H, W, 3].
                - mask_tensor (torch.Tensor or None): Normalized mask tensor [H, W].
                - first_dims (tuple): The updated or original reference dimensions.

        Raises:
            ValueError: If heterogeneous_dims is False and a dimension mismatch is found.
        """
        with archive.open(filename) as file_obj:
            try:
                img = Image.open(io.BytesIO(file_obj.read()))
                img.load()  # Force load the image data
            except OSError as e:
                logging.warning("Failed to load image %s: %s", filename, e)
                return None, None, first_dims

            if first_dims is None:
                first_dims = img.size
            elif img.size != first_dims:
                if not heterogeneous_dims:
                    raise ValueError(
                        f"Dimension mismatch: {filename} has dimensions "
                        f"{img.size}, expected {first_dims}."
                    )
                logging.warning(
                    "Skipping %s due to dimension mismatch: %s != %s",
                    filename,
                    img.size,
                    first_dims,
                )
                return None, None, first_dims

            rgb_img = img.convert("RGB")
            mask_img = img.convert("L")

            rgb_tensor = torch.from_numpy(np.array(rgb_img).astype(np.float32) / 255.0)
            mask_tensor = torch.from_numpy(np.array(mask_img).astype(np.float32) / 255.0)

            return rgb_tensor, mask_tensor, first_dims

    def load_from_zip(self, zip_file: str, heterogeneous_dimensions: bool):
        """
        The main entry point for the node. Loads and batches images from the ZIP.

        Args:
            zip_file (str): The filename of the ZIP archive located in the
                ComfyUI input directory.
            heterogeneous_dimensions (bool): If True, the node will skip images
                that do not match the dimensions of the first valid image found
                in the archive. If False, any mismatch will raise a ValueError.

        Returns:
            tuple: A tuple containing:
                - IMAGE (torch.Tensor): A batched tensor of shape [B, H, W, 3].
                - MASK (torch.Tensor): A batched tensor of shape [B, H, W].
                - COUNT (int): The number of images successfully loaded.

        Raises:
            FileNotFoundError: If the specified ZIP file does not exist.
            ValueError: If no valid images are found or if processing fails due
                to dimension mismatches (when heterogeneous_dimensions is False).
        """
        zip_path = os.path.join(folder_paths.get_input_directory(), zip_file)

        if not os.path.exists(zip_path):
            raise FileNotFoundError(f"ZIP file not found: {zip_path}")

        images = []
        masks = []

        with zipfile.ZipFile(zip_path, "r") as archive:
            filenames = [
                f
                for f in archive.namelist()
                if f.lower().endswith((".png", ".jpg", ".jpeg"))
                and not f.startswith("__MACOSX/")
                and not os.path.basename(f).startswith(".")
            ]

            filenames.sort()

            if not filenames:
                raise ValueError(f"No valid images found in the ZIP archive: {zip_file}")

            first_dimensions = None

            for filename in filenames:
                rgb_tensor, mask_tensor, first_dimensions = self._process_image_file(
                    archive, filename, first_dimensions, heterogeneous_dimensions
                )
                if rgb_tensor is not None and mask_tensor is not None:
                    images.append(rgb_tensor)
                    masks.append(mask_tensor)

        if not images:
            raise ValueError(f"No valid images could be processed from the ZIP archive: {zip_file}")

        batched_images = torch.stack(images)
        batched_masks = torch.stack(masks)

        return batched_images, batched_masks, len(images)