From b54b4329ca1e9d0beebdcd98a755ebf753e61e95 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Tue, 20 Jan 2026 00:20:21 +0100 Subject: [PATCH 1/2] Update fast_saver.py --- fast_saver.py | 103 ++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 82 insertions(+), 21 deletions(-) diff --git a/fast_saver.py b/fast_saver.py index 875fc3c..3569931 100644 --- a/fast_saver.py +++ b/fast_saver.py @@ -1,11 +1,12 @@ import os import torch import numpy as np -from PIL import Image +from PIL import Image, ExifTags from PIL.PngImagePlugin import PngInfo import concurrent.futures import re import time +import json class FastAbsoluteSaver: @classmethod @@ -15,9 +16,18 @@ class FastAbsoluteSaver: "images": ("IMAGE", ), "output_path": ("STRING", {"default": "D:\\Datasets\\Sharp_Output"}), "filename_prefix": ("STRING", {"default": "frame"}), - "metadata_key": ("STRING", {"default": "sharpness_score"}), - # NEW: Boolean Switch + + # --- FORMAT SWITCH --- + "save_format": (["png", "webp"], ), + + # --- COMMON OPTIONS --- "filename_with_score": ("BOOLEAN", {"default": False, "label": "Append Score to Filename"}), + "metadata_key": ("STRING", {"default": "sharpness_score"}), + + # --- WEBP SPECIFIC (-z 6 -q 100 -lossless) --- + "webp_lossless": ("BOOLEAN", {"default": True, "label": "WebP Lossless"}), + "webp_quality": ("INT", {"default": 100, "min": 0, "max": 100, "step": 1, "label": "WebP Quality (-q)"}), + "webp_method": ("INT", {"default": 4, "min": 0, "max": 6, "step": 1, "label": "WebP Compression (-z)"}), }, "optional": { "scores_info": ("STRING", {"forceInput": True}), @@ -30,17 +40,12 @@ class FastAbsoluteSaver: CATEGORY = "BetaHelper/IO" def parse_info(self, info_str, batch_size): - """ - Extracts both Frame Indices AND Scores. - """ if not info_str: return ([0]*batch_size, [0.0]*batch_size) matches = re.findall(r"F:(\d+).*?Score:\s*(\d+(\.\d+)?)", info_str) - frames = [] scores = [] - for m in matches: try: frames.append(int(m[0])) @@ -55,22 +60,68 @@ class FastAbsoluteSaver: return frames[:batch_size], scores[:batch_size] - def save_single_image(self, tensor_img, full_path, score, key_name): + def get_webp_exif(self, key, value): + """ + Creates a basic Exif header to store the score in UserComment. + ComfyUI standard metadata handling for WebP is complex, + so we use a simple JSON dump inside the UserComment tag (ID 0x9286). + """ + # Create a basic exif dict + exif_data = { + 0x9286: f"{key}: {value}".encode("utf-8") # UserComment + } + + # Convert to bytes manually to avoid requiring 'piexif' library + # This is a minimal TIFF header structure for Exif. + # If this is too hacky, we can just skip metadata for WebP, + # but this usually works for basic viewers. + + # ACTUALLY: Pillow's image.save(exif=...) expects raw bytes. + # Generating raw Exif bytes from scratch is error-prone. + # Simpler Strategy: We will create a fresh Image and modify its info. + + # Since generating raw Exif without a library is risky, + # we will skip internal metadata for WebP in this "No-Dependency" version + # and rely on the filename. + # *However*, if you strictly need it, we return None here and rely on filename + # unless you have 'piexif' installed. + return None + + def save_single_image(self, tensor_img, full_path, score, key_name, fmt, lossless, quality, method): try: array = 255. * tensor_img.cpu().numpy() img = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8)) - metadata = PngInfo() - metadata.add_text(key_name, str(score)) - metadata.add_text("software", "ComfyUI_Parallel_Node") - - img.save(full_path, pnginfo=metadata, compress_level=1) + if fmt == "png": + # PNG METADATA (Robust) + metadata = PngInfo() + metadata.add_text(key_name, str(score)) + metadata.add_text("software", "ComfyUI_Parallel_Node") + # PNG uses compress_level (0-9). Level 1 is fastest. + img.save(full_path, format="PNG", pnginfo=metadata, compress_level=1) + + elif fmt == "webp": + # WEBP SAVING + # Pillow options map directly to cwebp parameters: + # method=6 -> -z 6 (Slowest, best compression) + # quality=100 -> -q 100 + # lossless=True -> -lossless + + # Note: WebP metadata in Pillow is finicky. + # We save purely visual data here. + # The score is in the filename (if option selected). + img.save(full_path, format="WEBP", + lossless=lossless, + quality=quality, + method=method) + return True except Exception as e: print(f"xx- Error saving {full_path}: {e}") return False - def save_images_fast(self, images, output_path, filename_prefix, metadata_key, filename_with_score, scores_info=None): + def save_images_fast(self, images, output_path, filename_prefix, save_format, filename_with_score, metadata_key, + webp_lossless, webp_quality, webp_method, scores_info=None): output_path = output_path.strip('"') if not os.path.exists(output_path): @@ -82,7 +133,7 @@ class FastAbsoluteSaver: batch_size = len(images) frame_indices, scores_list = self.parse_info(scores_info, batch_size) - print(f"xx- FastSaver: Saving {batch_size} images to {output_path}...") + print(f"xx- FastSaver: Saving {batch_size} images ({save_format}) to {output_path}...") with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor: futures = [] @@ -92,21 +143,31 @@ class FastAbsoluteSaver: real_frame_num = frame_indices[i] current_score = scores_list[i] - # BASE NAME: frame_001450 base_name = f"{filename_prefix}_{real_frame_num:06d}" - # OPTION: Append Score -> frame_001450_1500 if filename_with_score: base_name += f"_{int(current_score)}" - # FALLBACK for missing data if real_frame_num == 0 and scores_info is None: base_name = f"{filename_prefix}_{int(time.time())}_{i:03d}" - fname = f"{base_name}.png" + # Append correct extension + ext = ".webp" if save_format == "webp" else ".png" + fname = f"{base_name}{ext}" full_path = os.path.join(output_path, fname) - futures.append(executor.submit(self.save_single_image, img_tensor, full_path, current_score, metadata_key)) + # Submit + futures.append(executor.submit( + self.save_single_image, + img_tensor, + full_path, + current_score, + metadata_key, + save_format, # fmt + webp_lossless, # lossless + webp_quality, # quality + webp_method # method + )) concurrent.futures.wait(futures) From f63b837a2c6e95c560b0f3d6a380f92aecd1237e Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Tue, 20 Jan 2026 00:21:42 +0100 Subject: [PATCH 2/2] Update fast_saver.py --- fast_saver.py | 74 +++++++++++++++------------------------------------ 1 file changed, 22 insertions(+), 52 deletions(-) diff --git a/fast_saver.py b/fast_saver.py index 3569931..138a409 100644 --- a/fast_saver.py +++ b/fast_saver.py @@ -1,12 +1,11 @@ import os import torch import numpy as np -from PIL import Image, ExifTags +from PIL import Image from PIL.PngImagePlugin import PngInfo import concurrent.futures import re import time -import json class FastAbsoluteSaver: @classmethod @@ -20,11 +19,14 @@ class FastAbsoluteSaver: # --- FORMAT SWITCH --- "save_format": (["png", "webp"], ), + # --- PERFORMANCE --- + "max_threads": ("INT", {"default": 0, "min": 0, "max": 128, "step": 1, "label": "Max Threads (0=Auto)"}), + # --- COMMON OPTIONS --- "filename_with_score": ("BOOLEAN", {"default": False, "label": "Append Score to Filename"}), "metadata_key": ("STRING", {"default": "sharpness_score"}), - # --- WEBP SPECIFIC (-z 6 -q 100 -lossless) --- + # --- WEBP SPECIFIC --- "webp_lossless": ("BOOLEAN", {"default": True, "label": "WebP Lossless"}), "webp_quality": ("INT", {"default": 100, "min": 0, "max": 100, "step": 1, "label": "WebP Quality (-q)"}), "webp_method": ("INT", {"default": 4, "min": 0, "max": 6, "step": 1, "label": "WebP Compression (-z)"}), @@ -60,56 +62,18 @@ class FastAbsoluteSaver: return frames[:batch_size], scores[:batch_size] - def get_webp_exif(self, key, value): - """ - Creates a basic Exif header to store the score in UserComment. - ComfyUI standard metadata handling for WebP is complex, - so we use a simple JSON dump inside the UserComment tag (ID 0x9286). - """ - # Create a basic exif dict - exif_data = { - 0x9286: f"{key}: {value}".encode("utf-8") # UserComment - } - - # Convert to bytes manually to avoid requiring 'piexif' library - # This is a minimal TIFF header structure for Exif. - # If this is too hacky, we can just skip metadata for WebP, - # but this usually works for basic viewers. - - # ACTUALLY: Pillow's image.save(exif=...) expects raw bytes. - # Generating raw Exif bytes from scratch is error-prone. - # Simpler Strategy: We will create a fresh Image and modify its info. - - # Since generating raw Exif without a library is risky, - # we will skip internal metadata for WebP in this "No-Dependency" version - # and rely on the filename. - # *However*, if you strictly need it, we return None here and rely on filename - # unless you have 'piexif' installed. - return None - def save_single_image(self, tensor_img, full_path, score, key_name, fmt, lossless, quality, method): try: array = 255. * tensor_img.cpu().numpy() img = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8)) if fmt == "png": - # PNG METADATA (Robust) metadata = PngInfo() metadata.add_text(key_name, str(score)) metadata.add_text("software", "ComfyUI_Parallel_Node") - # PNG uses compress_level (0-9). Level 1 is fastest. img.save(full_path, format="PNG", pnginfo=metadata, compress_level=1) elif fmt == "webp": - # WEBP SAVING - # Pillow options map directly to cwebp parameters: - # method=6 -> -z 6 (Slowest, best compression) - # quality=100 -> -q 100 - # lossless=True -> -lossless - - # Note: WebP metadata in Pillow is finicky. - # We save purely visual data here. - # The score is in the filename (if option selected). img.save(full_path, format="WEBP", lossless=lossless, quality=quality, @@ -120,8 +84,8 @@ class FastAbsoluteSaver: print(f"xx- Error saving {full_path}: {e}") return False - def save_images_fast(self, images, output_path, filename_prefix, save_format, filename_with_score, metadata_key, - webp_lossless, webp_quality, webp_method, scores_info=None): + def save_images_fast(self, images, output_path, filename_prefix, save_format, max_threads, + filename_with_score, metadata_key, webp_lossless, webp_quality, webp_method, scores_info=None): output_path = output_path.strip('"') if not os.path.exists(output_path): @@ -130,12 +94,20 @@ class FastAbsoluteSaver: except OSError: raise ValueError(f"Could not create directory: {output_path}") + # --- AUTO-SCALING LOGIC --- + if max_threads == 0: + # os.cpu_count() returns None on some rare systems, so we default to 4 just in case + cpu_cores = os.cpu_count() or 4 + # For WebP (CPU intensive), stick to core count. + # For PNG (Disk intensive), we could technically go higher, but core count is safe. + max_threads = cpu_cores + + print(f"xx- FastSaver: Using {max_threads} Threads for saving.") + batch_size = len(images) frame_indices, scores_list = self.parse_info(scores_info, batch_size) - print(f"xx- FastSaver: Saving {batch_size} images ({save_format}) to {output_path}...") - - with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor: + with concurrent.futures.ThreadPoolExecutor(max_workers=max_threads) as executor: futures = [] for i, img_tensor in enumerate(images): @@ -151,22 +123,20 @@ class FastAbsoluteSaver: if real_frame_num == 0 and scores_info is None: base_name = f"{filename_prefix}_{int(time.time())}_{i:03d}" - # Append correct extension ext = ".webp" if save_format == "webp" else ".png" fname = f"{base_name}{ext}" full_path = os.path.join(output_path, fname) - # Submit futures.append(executor.submit( self.save_single_image, img_tensor, full_path, current_score, metadata_key, - save_format, # fmt - webp_lossless, # lossless - webp_quality, # quality - webp_method # method + save_format, + webp_lossless, + webp_quality, + webp_method )) concurrent.futures.wait(futures)