commit 73217ca6320ee638dfbcb9e164b63a3f80cf34ef Author: hmwl <402245847@qq.com> Date: Thu Jan 16 09:50:39 2025 +0800 first commit diff --git a/Example.json b/Example.json new file mode 100644 index 0000000..eba8c75 --- /dev/null +++ b/Example.json @@ -0,0 +1,220 @@ +{ + "last_node_id": 15, + "last_link_id": 14, + "nodes": [ + { + "id": 15, + "type": "CleanFolders", + "pos": [ + 1179.648193359375, + 803.1049194335938 + ], + "size": [ + 210, + 54 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "output_path", + "type": "STRING", + "link": 13, + "widget": { + "name": "output_path" + }, + "label": "output_path" + }, + { + "name": "input_path", + "type": "STRING", + "link": 14, + "widget": { + "name": "input_path" + }, + "label": "input_path" + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "CleanFolders" + }, + "widgets_values": [ + "", + "" + ] + }, + { + "id": 11, + "type": "CompressImages", + "pos": [ + 894.0198974609375, + 745.897705078125 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 9, + "label": "images" + } + ], + "outputs": [ + { + "name": "temp_dir", + "type": "STRING", + "links": [ + 13 + ], + "slot_index": 0, + "label": "temp_dir" + } + ], + "properties": { + "Node name for S&R": "CompressImages" + }, + "widgets_values": [ + "ComfyUI" + ] + }, + { + "id": 13, + "type": "ImageInvert", + "pos": [ + 698.370849609375, + 745.5828247070312 + ], + "size": [ + 140, + 26 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 8, + "label": "image" + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 9 + ], + "slot_index": 0, + "label": "IMAGE" + } + ], + "properties": { + "Node name for S&R": "ImageInvert" + }, + "widgets_values": [] + }, + { + "id": 9, + "type": "UnzipToInput", + "pos": [ + 430.0738220214844, + 798.1468505859375 + ], + "size": [ + 210, + 78 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 8 + ], + "slot_index": 0, + "shape": 6, + "label": "images" + }, + { + "name": "folder", + "type": "STRING", + "links": [ + 14 + ], + "slot_index": 1, + "label": "folder" + } + ], + "properties": { + "Node name for S&R": "UnzipToInput" + }, + "widgets_values": [ + "input/file.zip" + ] + } + ], + "links": [ + [ + 8, + 9, + 0, + 13, + 0, + "IMAGE" + ], + [ + 9, + 13, + 0, + 11, + 0, + "IMAGE" + ], + [ + 13, + 11, + 0, + 15, + 0, + "STRING" + ], + [ + 14, + 9, + 1, + 15, + 1, + "STRING" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.8769226950000008, + "offset": [ + -24.799331551459204, + -332.5597278000348 + ] + }, + "groupNodes": {}, + "node_versions": { + "comfy-core": "0.3.10" + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..0feeb6a --- /dev/null +++ b/__init__.py @@ -0,0 +1,17 @@ +from .unzip import UnzipToInput +from .compress import CompressImages +from .clean import CleanFolders + +NODE_CLASS_MAPPINGS = { + "UnzipToInput": UnzipToInput, + "CompressImages": CompressImages, + "CleanFolders": CleanFolders +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "UnzipToInput": "Unzip To Input", + "CompressImages": "Zip Compress", + "CleanFolders": "Clean Folders" +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/clean.py b/clean.py new file mode 100644 index 0000000..454b323 --- /dev/null +++ b/clean.py @@ -0,0 +1,61 @@ +import os +import shutil +import folder_paths + +class CleanFolders: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "input_path": ("STRING", {"default": ""}), + "output_path": ("STRING", {"default": ""}) + }} + + RETURN_TYPES = () + FUNCTION = "clean" + CATEGORY = "comfyui_zip" + OUTPUT_NODE = True + + def clean(self, input_path, output_path): + result_msg = [] + + # 获取基础目录 + input_base = folder_paths.get_input_directory() + output_base = folder_paths.get_output_directory() + + # 确保路径是字符串类型，如果是元组则取第一个元素 + input_path = input_path[0] if isinstance(input_path, tuple) else str(input_path) if input_path else "" + output_path = output_path[0] if isinstance(output_path, tuple) else str(output_path) if output_path else "" + + # 只有当提供了路径时才进行删除操作 + if input_path: + input_full_path = os.path.join(input_base, input_path) + if os.path.exists(input_full_path): + try: + shutil.rmtree(input_full_path) + result_msg.append(f"已清理输入目录: {input_full_path}") + except Exception as e: + result_msg.append(f"清理输入目录失败: {str(e)}") + else: + result_msg.append(f"输入目录不存在: {input_full_path}") + else: + result_msg.append("未提供输入路径，跳过清理") + + # 只有当提供了路径时才进行删除操作 + if output_path: + output_full_path = os.path.join(output_base, output_path) + if os.path.exists(output_full_path): + try: + shutil.rmtree(output_full_path) + except Exception as e: + pass + + return ("",) + +# 注册节点 +NODE_CLASS_MAPPINGS = { + "CleanFolders": CleanFolders +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "CleanFolders": "清理文件夹" +} diff --git a/compress.py b/compress.py new file mode 100644 index 0000000..3f8393c --- /dev/null +++ b/compress.py @@ -0,0 +1,118 @@ +import os +import zipfile +import torch +import numpy as np +from PIL import Image +import folder_paths +import tempfile +import shutil +import matplotlib.pyplot as plt +import time +import random + +class CompressImages: + _current_id = None + _temp_dirs = {} + _image_counts = {} + _zip_filename = {} # 新增：记录每个工作流的zip文件名 + + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "filename_prefix": ("STRING", {"default": "ComfyUI"}) + }, + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + "unique_id": "UNIQUE_ID", + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("temp_dir",) + FUNCTION = "compress_images" + OUTPUT_NODE = True + CATEGORY = "comfyui_zip" + + def compress_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, unique_id=None): + try: + # 确保临时目录存在 + if unique_id not in CompressImages._temp_dirs or not os.path.exists(CompressImages._temp_dirs[unique_id]): + folder_name = f"{filename_prefix}{self.prefix_append}_{unique_id}" + CompressImages._temp_dirs[unique_id] = os.path.join(self.output_dir, folder_name) + CompressImages._image_counts[unique_id] = 0 + timestamp = int(time.time()) + CompressImages._zip_filename[unique_id] = f'{filename_prefix}_{timestamp}.zip' + + # 确保目录存在 + os.makedirs(CompressImages._temp_dirs[unique_id], exist_ok=True) + + temp_dir = CompressImages._temp_dirs[unique_id] + zip_path = os.path.abspath(os.path.join(self.output_dir, CompressImages._zip_filename[unique_id])) + + # 追加模式打开zip文件 + zip_mode = 'a' if os.path.exists(zip_path) else 'w' + with zipfile.ZipFile(zip_path, zip_mode) as zipf: + for image in images: + if isinstance(image, torch.Tensor): + img_array = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(img_array, 0, 255).astype(np.uint8)) + elif isinstance(image, Image.Image): + img = image + else: + raise ValueError(f"不支持的图片格式: {type(image)}") + + temp_image_path = os.path.join(temp_dir, f'temp_image_{CompressImages._image_counts[unique_id]:05d}.png') + img.save(temp_image_path) + arcname = f'image_{CompressImages._image_counts[unique_id]:05d}.png' + zipf.write(temp_image_path, arcname) + os.remove(temp_image_path) + CompressImages._image_counts[unique_id] += 1 + + folder_name = os.path.basename(CompressImages._temp_dirs[unique_id]) + + # 正确的返回格式 + results = (folder_name,) # 这是 RETURN_TYPES 定义的返回值 + + return {"ui": {"file": [CompressImages._zip_filename[unique_id]]}, "result": (results,)} + + except Exception as e: + print(f"发生错误: {str(e)}") + self._cleanup(unique_id) + raise e + +class ZipFiles: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "images": ("IMAGE",), + "zip_name": ("STRING", {"default": "images.zip"}), + "save_path": ("STRING", {"default": "ComfyUI/output"}), + }, + "hidden": {"unique_id": "UNIQUE_ID"}, + } + + RETURN_TYPES = ("STRING", "STRING",) + RETURN_NAMES = ("zip_path", "temp_dir",) + FUNCTION = "zip_files" + + def zip_files(self, images, zip_name, save_path, unique_id): + + temp_dir = tempfile.gettempdir() + + return (output_path, temp_dir) + +NODE_CLASS_MAPPINGS = { + "CompressImages": CompressImages +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "CompressImages": "Zip Compress" +} diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..f316d6c --- /dev/null +++ b/readme.md @@ -0,0 +1,58 @@ +## ComfyUI 压缩包图片处理 +## ComfyUI zip package image processing + +### 节点介绍 +### Node Introduction + +#### 1. Zip Compress (压缩图片) +- 功能：将生成的图片打包成zip压缩文件 +- Function: Package generated images into a zip file +- 输入参数 (Input Parameters): + - images: 图片输入 (Image input) + - filename_prefix: 文件名前缀，默认为"ComfyUI" (Filename prefix, default is "ComfyUI") +- 输出 (Output): + - temp_dir: 临时目录路径 (Temporary directory path) + - 生成的zip文件将保存在ComfyUI的输出目录中 (Generated zip file will be saved in ComfyUI's output directory) + +#### 2. UnzipToInput (解压到输入) +- 功能：解压zip文件并将图片加载为可处理的格式 +- Function: Extract zip file and load images into processable format +- 输入参数 (Input Parameters): + - zip_path: zip文件路径，支持本地路径或URL (Zip file path, supports local path or URL) +- 输出 (Output): + - images: 解压出的图片列表 (List of extracted images) + - folder: 解压目录名称 (Extraction directory name) + +> **注意**: 压缩包仅支持图片文件的zip，不支持其他类型文件和文件夹 +> **Note**: Only supports zip files containing images, other file types and folders are not supported + +#### 3. 清理文件夹 (Clean Folders) +- 功能：清理临时文件夹和输出文件夹 +- Function: Clean temporary and output folders +- 输入参数 (Input Parameters): + - input_path: 需要清理的输入路径 (Input path to clean) + - output_path: 需要清理的输出路径 (Output path to clean) +- 输出 (Output): + - 无返回值，直接执行清理操作 (No return value, directly performs cleaning operation) + +### 使用说明 +### Usage Instructions + +1. 压缩图片 (Compress Images): + - 将图片节点连接到Zip Compress节点 + - 设置文件名前缀 + - 运行后将在输出目录生成zip文件 + +2. 解压图片 (Extract Images): + - 在UnzipToInput节点中输入zip文件路径 + - 支持本地zip文件或在线zip文件URL + - 解压后的图片可直接用于后续处理 + +3. 清理文件夹 (Clean Folders): + - 用于清理工作流程中产生的临时文件 + - 可选择性地清理输入或输出目录 + - 建议在工作流程结束时使用 + +### 测试JSON文件 +项目中有一个Example.json文件，可以作为使用的参考 +There is an Example.json file in the project that you can refer to for usage. diff --git a/unzip.py b/unzip.py new file mode 100644 index 0000000..72c4a7b --- /dev/null +++ b/unzip.py @@ -0,0 +1,103 @@ +import os +import time +import torch +import urllib.request +from urllib.parse import urlparse +import folder_paths +import numpy as np +from PIL import Image +import shutil + +class UnzipToInput: + def __init__(self): + self.output_dir = folder_paths.get_input_directory() + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "zip_path": ("STRING", {"default": ""}), + } + } + + RETURN_TYPES = ("IMAGE", "STRING",) + RETURN_NAMES = ("images", "folder",) + OUTPUT_IS_LIST = (True, False,) + FUNCTION = "unzip_process" + CATEGORY = "comfyui_zip" + + def download_file(self, url, save_path): + urllib.request.urlretrieve(url, save_path) + + def is_image_file(self, filename): + image_extensions = {'.png', '.jpg', '.jpeg', '.webp', '.bmp'} + return os.path.splitext(filename.lower())[1] in image_extensions + + def unzip_process(self, zip_path): + timestamp = str(int(time.time())) + extract_dir = os.path.join(self.output_dir, f"unzip_{timestamp}") + os.makedirs(extract_dir, exist_ok=True) + + # 处理URL或本地路径 + if zip_path.startswith(('http://', 'https://')): + parsed_url = urlparse(zip_path) + zip_filename = os.path.basename(parsed_url.path) + temp_zip = os.path.join(extract_dir, zip_filename) + self.download_file(zip_path, temp_zip) + zip_path = temp_zip + + # 解压文件 + import zipfile + with zipfile.ZipFile(zip_path, 'r') as zip_ref: + # 获取所有文件列表 + name_list = zip_ref.namelist() + + # 过滤掉 __MACOSX 文件夹 + filtered_names = [name for name in name_list + if not name.startswith('__MACOSX/') + and not name.startswith('._')] + + # 解压文件，处理中文编码 + for name in filtered_names: + try: + # 尝试使用 cp437 解码，然后用 utf-8 编码 + unicode_name = name.encode('cp437').decode('utf-8') + except: + # 如果失败，直接使用原名称 + unicode_name = name + + # 获取解压后的完整路径 + extracted_path = os.path.join(extract_dir, unicode_name) + + # 确保目标目录存在 + os.makedirs(os.path.dirname(extracted_path), exist_ok=True) + + # 提取文件 + with zip_ref.open(name) as source, open(extracted_path, 'wb') as target: + shutil.copyfileobj(source, target) + + # 收集所有图片文件 + image_files = [] + for root, _, files in os.walk(extract_dir): + for file in files: + if self.is_image_file(file): + image_files.append(os.path.join(root, file)) + + if not image_files: + raise ValueError("No image files found in ZIP archive") + + # 读取所有图片并转换为tensor + images_list = [] + for img_path in image_files: + img = Image.open(img_path) + if img.mode != 'RGB': + img = img.convert('RGB') + img_tensor = np.array(img).astype(np.float32) / 255.0 + img_tensor = torch.from_numpy(img_tensor) + img_tensor = img_tensor.unsqueeze(0) + images_list.append(img_tensor) + + # 返回图片列表和文件夹名称 + folder_name = os.path.basename(extract_dir) + + return (images_list, folder_name)