diff --git a/zip_batch_loader.py b/zip_batch_loader.py index f36ba3b..ed5e182 100644 --- a/zip_batch_loader.py +++ b/zip_batch_loader.py @@ -17,12 +17,33 @@ import folder_paths # pylint: disable=import-error # noinspection PyPep8Naming class ZipBatchLoader: """ - A custom node to load a batch of images directly from a ZIP archive. + 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 types for the ComfyUI node.""" + """ + 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() @@ -47,7 +68,28 @@ class ZipBatchLoader: CATEGORY = "image" def _process_image_file(self, archive, filename, first_dims, heterogeneous_dims): - """Helper to process a single image from the archive.""" + """ + 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())) @@ -81,7 +123,27 @@ class ZipBatchLoader: return rgb_tensor, mask_tensor, first_dims def load_from_zip(self, zip_file: str, heterogeneous_dimensions: bool): - """Loads images from the given zip file and returns batched tensors.""" + """ + 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):